Cortex Core vs LangGraph

00 Overview

One app, built twice. Read this first, then step through the files in order.

Cortex is a chatbot with two jobs: an agent that remembers past conversations and can search the web, and a RAG pipeline that answers questions from uploaded PDFs. cortex-core builds everything by hand on the raw OpenAI, transformers and chromadb libraries. cortex-langgraph rebuilds the same app on LangChain components and LangGraph graphs, keeping prompts, defaults and behavior the same so the two can be compared fairly.

The pattern to look for

Almost every change is one of two moves. LangChain replaces a hand-written interface (model, embeddings, vector store, prompt, splitter) with a standard one that every provider implements. LangGraph replaces hand-written control flow (the tool loop, the RAG steps, the fallback policy) with a graph of nodes and edges you can see, bound and trace.

What replaced what

cortex-core (hand-written)		cortex-langgraph	Tab
`BaseLLM` + `OpenAIModel` / `LocalHFModel`	→	`BaseChatModel` (`ChatOpenAI`, `AzureChatOpenAI`, `ChatHuggingFace`)	model.py
`Embedder` + hand-written retry loop	→	`Embeddings` + `CacheBackedEmbeddings`	embedding.py
`PromptBuilder` f-strings	→	`ChatPromptTemplate` / `PromptTemplate`	prompt_builder.py
`while True:` tool loop	→	`StateGraph` + `ToolNode` + `tools_condition`	agentic_chatbot.py
`TOOLS` JSON schemas + `if/elif` dispatch	→	`@tool` functions	tools.py
`RAGPipeline.answer()`, three calls in a row	→	RAG `StateGraph` with a policy branch	retrieval.py
`select_context()` fallback policy	→	Conditional edge `_route_after_rerank()`	orchestrator.py
`VectorStoreBase` + two backends	→	`VectorStore` (`Chroma`, `AzureSearch`)	vector_store.py
Fixed-window `Chunker`	→	`RecursiveCharacterTextSplitter`	chunk.py
Hand-walked Docling model	→	`DoclingLoader`	pdf_digest.py
Manual torch cross-encoder	→	`BaseCrossEncoder` over sentence-transformers	cross_encoder.py
`RetrievedChunk` / `ReRankedChunk` dataclasses	→	`Document` + `metadata`	interface.py
Token usage never reported	→	`AIMessage.usage_metadata` on every call	chatbot/__init__.py

How to read each tab

Each tab starts with a short note: what the file does, how each project does it, and why the change matters. Below it is an aligned side-by-side diff: cortex-core on the left, cortex-langgraph on the right. Red rows exist only on the left or were rewritten, green rows are new on the right, and white rows are identical. Turn on Hide unchanged lines to see only the differences.

17

files changed

18

files identical

3

files only in cortex-core

1

file only in cortex-langgraph

Unchanged between the two

The whole API layer (FastAPI routes, controller, rate limiter, metrics), the data models, cost calculation, web search and the Redis / SQLite memory are byte-for-byte identical. The framework swap only touched the model, retrieval and agent code.

scripts/index_cli.py
src/api/controller.py
src/api/main.py
src/api/metrics.py
src/api/ratelimiter.py
src/api/router.py
src/api/service.py
src/business/core/cost.py
src/business/core/live_data.py
src/business/rag/re_ranker/__init__.py
src/business/rag/re_ranker/config.py
src/config/config.yml
src/database/database.py
src/database/dto.py
src/memory/chat_history_manager.py
src/memory/redis_memory.py
src/memory/responsecache.py
src/utils/config.py

Also only in cortex-core: re_ranker/gating.py, an empty placeholder file.

New dependencies in cortex-langgraph

langchainlangchain-corelanggraphlanggraph-prebuiltlangchain-openailangchain-huggingfacelangchain-chromalangchain-azure-ailangchain-doclinglangchain-text-splitterslangchain-classiclangsmith

01 model.py

src/business/core/model.py · 183 → 147 lines · 258 rows changed

How the app talks to a chat model. cortex-core wrote its own model interface; LangChain already ships one that every provider implements.

cortex-core · hand-written

A hand-written BaseLLM abstract class with one method, generate(question, context).
Two implementations: OpenAIModel wraps the raw OpenAI client, LocalHFModel drives transformers + torch directly (tokenizer, model.generate, decode, strip the prompt).
build_azure_openai_client() returns a raw client so the agent can do tool calling separately.

cortex-langgraph · LangChain / LangGraph

No interface of its own: create_llm() returns a LangChain BaseChatModel (ChatOpenAI, AzureChatOpenAI or ChatHuggingFace).
Every provider is called the same way: llm.invoke(messages) returns an AIMessage.
Built-in retries (max_retries=5) and tool calling (.bind_tools()) on the same object.
model_label() reads the model name back for metrics, whichever provider it is.

Why it matters: Every response now carries token usage (AIMessage.usage_metadata). cortex-core never reported tokens, so its cost tracking could never run.

BaseChatModelChatOpenAIAzureChatOpenAIChatHuggingFace

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/core/model.py		cortex-langgraph src/business/core/model.py
1	#Here is the openai LLM model implementation	1	"""Chat model factory — LangChain edition.
		2
		3	In cortex-core this file held a hand-written `BaseLLM` interface with two
		4	implementations (`OpenAIModel` wrapping the raw OpenAI client, `LocalHFModel`
		5	wrapping transformers + torch directly), each with its own `generate()`.
		6
		7	LangChain already defines that interface: `BaseChatModel`. Every provider
		8	package ships an implementation of it, so the three branches below just pick
		9	one. Everything downstream (the RAG graph, the agent graph) only ever sees a
		10	`BaseChatModel` and calls `.invoke()` / `.bind_tools()` on it — it never
		11	knows which provider is behind it.
		12
		13	What LangChain gives us for free compared to cortex-core:
		14	- One call shape for every provider: `llm.invoke(messages) -> AIMessage`.
		15	- Token usage on every response (`AIMessage.usage_metadata`), which cortex-core
		16	never surfaced — that is what finally makes cost tracking work end-to-end.
		17	- Tool calling (`.bind_tools()`) on the same object, so the agent no longer
		18	needs a separate raw OpenAI client.
		19	- Built-in retries (`max_retries`).
		20	"""
		21
		22	from __future__ import annotations
2		23
3	import os	24	import os
4	from transformers import AutoModelForCausalLM, AutoTokenizer	25	from typing import Optional
5	from abc import ABC, abstractmethod	26
6	import torch
7	from .prompt_builder import PromptBuilder
8	from openai import OpenAI, AzureOpenAI
9	from dotenv import load_dotenv	27	from dotenv import load_dotenv
10	from typing import Optional	28	from langchain_core.language_models import BaseChatModel
		29
11	load_dotenv()	30	load_dotenv()
12
13	# Base LLM Interface
14	class BaseLLM(ABC):
15	@abstractmethod
16	def generate(self, question: str, context: list[str]) -> str:
17	pass
18	# abstractmethod used when you have 2 or more models that you will decide to use which one to use at runtime
19
20	#HuggingFace Local LLM Implementation
21	class LocalHFModel(BaseLLM):
22	"""
23	Wrapper for openai LLMs to generate responses based on prompts.
24	Local Hugging Face Model Implementation
25	suitable for on-prem or offline inference
26	"""
27	def __init__(self, model_name: str, system_prompt: str, max_input_tokens: int = 2048, max_output_tokens: int = 512):
28
29	self.tokenizer = AutoTokenizer.from_pretrained(model_name)
30	self.model = AutoModelForCausalLM.from_pretrained(model_name)
31	self.prompt_builder = PromptBuilder(system_prompt)
32	self.max_input_tokens = max_input_tokens
33	self.max_output_tokens = max_output_tokens
34
35	if self.tokenizer.pad_token is None:
36	self.tokenizer.pad_token = self.tokenizer.eos_token
37	"""
38	Simple meaning:
39	If the tokenizer does not define a padding token,
40	the end-of-sequence (EOS) token is used as the padding token.
41
42	Explanation:
43	The pad_token is required to make all sequences in a batch the same length.
44
45	Example:
46	"Hello"
47	"Hello how are you"
48
49	Converted to tensors:
50	[Hello, <PAD>, <PAD>]
51	[Hello, how, are, you]
52
53	Without a pad_token, batch processing will fail and raise runtime errors.
54
55	The eos_token (End Of Sequence) indicates the end of a sentence.
56	Many autoregressive models (e.g., GPT-style):
57	- Are not designed for batch padding
58	- Do not define a pad_token (pad_token = None)
59	- Only provide an eos_token
60
61	Using EOS as PAD:
62	- Prevents crashes and runtime errors
63	- Does not negatively affect attention mechanisms
64	- Is a common and safe practice during inference
65	"""
66	def generate(self, question: str, context: list[str]) -> str:
67	"""
68	Generates a response from the LLM based on the provided question and context.
69
70	Args:
71	question (str): The user's question or input.
72	context (list): Relevant context passages retrieved from the knowledge base.
73	max_length (int): Maximum length of the generated response.
74
75	Returns:
76	str: The generated response from the LLM.
77	"""
78	prompt = self.prompt_builder.build_prompt(question, context)
79	inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True,max_length=self.max_input_tokens)
80
81	with torch.no_grad():
82	outputs = self.model.generate(**inputs, max_length=self.max_input_tokens, pad_token_id=self.tokenizer.eos_token_id)
83
84	generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
85	# Extract only the part after the prompt
86	return generated_text[len(prompt):].strip()
87
88	#OPENAI LLM Implementation
89
90	class OpenAIModel(BaseLLM):
91	def __init__(self, client: OpenAI, model_name: str, system_prompt:str, temperature: float = 0.7, max_tokens: int = 512):
92	self.client = client
93	self.model_name = model_name
94	self.temperature = temperature
95	self.max_tokens = max_tokens
96	self.prompt_builder = PromptBuilder(system_prompt)
97	def generate(self, question: str, context: list[str]) -> str:
98	"""
99	Generates a response from the OpenAI LLM based on the provided question and context.
100	Args:
101	question (str): The user's question or input.
102	context (list): Relevant context passages retrieved from the knowledge base.
103	Returns:
104	str: The generated response from the LLM.
105	"""
106	messages = self.prompt_builder.build_messages(question, context)
107	response = self.client.chat.completions.create(
108	model=self.model_name,
109	messages=messages,
110	max_tokens=self.max_tokens,
111	temperature=self.temperature
112	)
113	return response.choices[0].message.content.strip()
114		31
115		32
116	# ── Provider selection ────────────────────────────────────────────────────	33	# ── Provider selection ────────────────────────────────────────────────────
⋯ 1 unchanged line hidden
117	#	34	#
118	# LLM_PROVIDER (env var) picks the chat model implementation. On-prem/local	35	# LLM_PROVIDER (env var) picks the chat model implementation. On-prem/local
119	# deployments don't need to set anything — the default is today's behavior.	36	# deployments don't need to set anything — the default is today's behavior.
120	#	37	#
121	# openai (default) — OpenAI's cloud API. Current on-prem/as-is setup.	38	# openai (default) — ChatOpenAI, OpenAI's cloud API.
122	# huggingface — fully local inference, no network calls at generation time.	39	# huggingface — ChatHuggingFace over a local transformers pipeline;
123	# azure_openai — Azure OpenAI Service. Requires AZURE_OPENAI_ENDPOINT,	40	# no network calls at generation time.
		41	# azure_openai — AzureChatOpenAI. Requires AZURE_OPENAI_ENDPOINT,
124	# AZURE_OPENAI_API_KEY, AZURE_OPENAI_CHAT_DEPLOYMENT_NAME	42	# AZURE_OPENAI_API_KEY, AZURE_OPENAI_CHAT_DEPLOYMENT_NAME
125	# (see .env for the full list).	43	# (see .env for the full list).
126		44
127	def build_azure_openai_client(api_key: Optional[str] = None) -> AzureOpenAI:	45	def resolve_provider(provider: Optional[str] = None) -> str:
128	"""Shared AzureOpenAI client builder — used by both create_llm() and any	46	"""Normalise the provider name the same way everywhere (factory, agent, cost)."""
129	caller that needs a raw client (e.g. the agentic tool-calling loop)."""	47	return (provider or os.getenv("LLM_PROVIDER", "openai")).strip().lower()
		48
		49
		50	def create_llm(
		51	provider: Optional[str] = None,
		52	*,
		53	model_name: Optional[str] = None,
		54	api_key: Optional[str] = None,
		55	temperature: Optional[float] = None,
		56	max_tokens: Optional[int] = None,
		57	) -> BaseChatModel:
		58	"""Factory for the chat model, selected by LLM_PROVIDER or the `provider` arg.
		59
		60	`temperature` / `max_tokens` are only forwarded when given, so callers that
		61	don't care (the agent) get the provider's defaults, exactly as cortex-core's
		62	raw `chat.completions.create()` call did.
		63	"""
		64	provider = resolve_provider(provider)
		65	sampling = {}
		66	if temperature is not None:
		67	sampling["temperature"] = temperature
		68	if max_tokens is not None:
		69	sampling["max_tokens"] = max_tokens
		70
		71	if provider == "openai":
		72	from langchain_openai import ChatOpenAI
		73
		74	resolved_key = api_key or os.getenv("OPENAI_API_KEY")
		75	if not resolved_key:
		76	raise RuntimeError("OPENAI_API_KEY must be set for LLM_PROVIDER=openai")
		77	return ChatOpenAI(
		78	model=model_name or os.getenv("OPENAI_MODEL_NAME", "gpt-4o"),
		79	api_key=resolved_key,
		80	max_retries=5,
		81	**sampling,
		82	)
		83
		84	if provider == "huggingface":
		85	# Imported lazily: transformers + torch are heavy, and only this
		86	# provider needs them at generation time.
		87	from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline
		88
		89	pipeline = HuggingFacePipeline.from_model_id(
		90	model_id=model_name or os.getenv("HF_MODEL_NAME", "mistral-7b-instruct-v0.1"),
		91	task="text-generation",
		92	pipeline_kwargs={"max_new_tokens": max_tokens or 512, "return_full_text": False},
		93	)
		94	return ChatHuggingFace(llm=pipeline)
		95
		96	if provider == "azure_openai":
		97	from langchain_openai import AzureChatOpenAI
		98
		99	deployment = model_name or os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME")
		100	if not deployment:
		101	raise RuntimeError(
		102	"LLM_PROVIDER=azure_openai requires AZURE_OPENAI_CHAT_DEPLOYMENT_NAME "
		103	"(the deployment name you gave the chat model in Azure — not the "
		104	"underlying model name)."
		105	)
		106	endpoint, resolved_key, api_version = azure_openai_settings(api_key)
		107	return AzureChatOpenAI(
		108	azure_endpoint=endpoint,
		109	api_key=resolved_key,
		110	api_version=api_version,
		111	azure_deployment=deployment,
		112	max_retries=5,
		113	**sampling,
		114	)
		115
		116	raise ValueError(
		117	f"Unknown LLM_PROVIDER={provider!r}. Supported: openai, huggingface, azure_openai."
		118	)
		119
		120
		121	def azure_openai_settings(api_key: Optional[str] = None) -> tuple[str, str, str]:
		122	"""Shared Azure OpenAI connection settings — used by both create_llm() and
		123	create_embedder(), so the two can never disagree about which endpoint,
		124	key and API version they talk to."""
130	endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")	125	endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
131	resolved_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")	126	resolved_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
132	api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2024-10-21")	127	api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2024-10-21")
⋯ 2 unchanged lines hidden
133	if not endpoint or not resolved_key:	128	if not endpoint or not resolved_key:
134	raise RuntimeError(	129	raise RuntimeError(
135	"AZURE_OPENAI_ENDPOINT and AZURE_OPENAI_API_KEY must be set to use "	130	"AZURE_OPENAI_ENDPOINT and AZURE_OPENAI_API_KEY must be set to use "
136	"an azure_openai provider."	131	"an azure_openai provider."
137	)	132	)
138	return AzureOpenAI(azure_endpoint=endpoint, api_key=resolved_key, api_version=api_version)	133	return endpoint, resolved_key, api_version
139		134
140		135
141	def create_llm(	136	def model_label(llm: BaseChatModel) -> str:
142	provider: Optional[str] = None,	137	"""Best-effort model name for metrics and the API response.
143	*,
144	model_name: Optional[str] = None,
145	system_prompt: str = "",
146	api_key: Optional[str] = None,
147	) -> BaseLLM:
148	"""Factory for the chat LLM, selected by LLM_PROVIDER or the `provider` arg."""
149	provider = (provider or os.getenv("LLM_PROVIDER", "openai")).strip().lower()
150		138
151	if provider == "openai":	139	ChatOpenAI exposes `model_name`, AzureChatOpenAI `deployment_name`, and
152	resolved_key = api_key or os.getenv("OPENAI_API_KEY")	140	ChatHuggingFace `model_id`; fall back to the class name for anything else
153	if not resolved_key:	141	(including the fake models used in tests).
154	raise RuntimeError("OPENAI_API_KEY must be set for LLM_PROVIDER=openai")	142	"""
155	return OpenAIModel(	143	for attr in ("deployment_name", "model_name", "model_id", "model"):
156	client=OpenAI(api_key=resolved_key),	144	value = getattr(llm, attr, None)
157	model_name=model_name or os.getenv("OPENAI_MODEL_NAME", "gpt-4o"),	145	if isinstance(value, str) and value:
158	system_prompt=system_prompt,	146	return value
159	)	147	return type(llm).__name__
160
161	if provider == "huggingface":
162	return LocalHFModel(
163	model_name=model_name or os.getenv("HF_MODEL_NAME", "mistral-7b-instruct-v0.1"),
164	system_prompt=system_prompt,
165	)
166
167	if provider == "azure_openai":
168	deployment = model_name or os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME")
169	if not deployment:
170	raise RuntimeError(
171	"LLM_PROVIDER=azure_openai requires AZURE_OPENAI_CHAT_DEPLOYMENT_NAME "
172	"(the deployment name you gave the chat model in Azure — not the "
173	"underlying model name)."
174	)
175	return OpenAIModel(
176	client=build_azure_openai_client(api_key),
177	model_name=deployment,
178	system_prompt=system_prompt,
179	)
180
181	raise ValueError(
182	f"Unknown LLM_PROVIDER={provider!r}. Supported: openai, huggingface, azure_openai."
183	)

02 embedding.py

src/business/core/embedding.py · 112 → 120 lines · 143 rows changed

Turning text into vectors. Same two methods on both sides; LangChain removes the hand-written retry loop and adds a cache.

cortex-core · hand-written

A hand-written Embedder abstract class with embed_documents() and embed_query().
OpenAIEmbedder calls client.embeddings.create() and retries 500 errors with its own exponential-backoff loop.
Azure works by injecting an AzureOpenAI client into the same class.

cortex-langgraph · LangChain / LangGraph

Returns LangChain's Embeddings interface: OpenAIEmbeddings or AzureOpenAIEmbeddings, retries built in.
Optional on-disk cache (CacheBackedEmbeddings + LocalFileStore), keyed by a SHA-256 hash of the text and namespaced by provider and model.
Vector stores take the embedder directly, so the stores embed text themselves.

Why it matters: Re-indexing the same PDF re-embeds identical chunks. With the cache those come from disk instead of a paid API call.

EmbeddingsOpenAIEmbeddingsCacheBackedEmbeddings

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/core/embedding.py		cortex-langgraph src/business/core/embedding.py
1	"""Embedding interfaces and OpenAI implementation."""	1	"""Embedding factory — LangChain edition.
		2
		3	cortex-core defined its own `Embedder` ABC and an `OpenAIEmbedder` with a
		4	hand-rolled exponential-backoff retry loop. LangChain's `Embeddings` interface
		5	has the same two methods (`embed_documents`, `embed_query`), and
		6	`OpenAIEmbeddings` / `AzureOpenAIEmbeddings` implement it with retries built in.
		7
		8	Because every LangChain vector store takes an `Embeddings` object directly,
		9	the vector stores now embed text themselves — callers hand them strings /
		10	Documents instead of pre-computed vectors.
		11
		12	New compared to cortex-core: an optional on-disk embedding cache
		13	(`CacheBackedEmbeddings`). Re-indexing the same PDF re-embeds identical chunks;
		14	with the cache those come from disk instead of a paid API call.
		15	"""
2		16
3	from __future__ import annotations	17	from __future__ import annotations
4		18
5	import os	19	import os
6	import time	20	import re
7	from abc import ABC, abstractmethod	21	from typing import Optional
8	from typing import List, Optional
9		22
10	from openai import OpenAI	23	from langchain_core.embeddings import Embeddings
11	from openai import InternalServerError
12		24
13		25	from .model import azure_openai_settings
14	class Embedder(ABC):
15	@abstractmethod
16	def embed_documents(self, texts: List[str]) -> List[List[float]]:
17	...
18
19	@abstractmethod
20	def embed_query(self, text: str) -> List[float]:
21	...
22
23	def embed(self, text: str) -> List[float]:
24	"""
25	Alias for embed_query() for backward compatibility with LongTermMemory.
26	"""
27	return self.embed_query(text)
28
29
30	class OpenAIEmbedder(Embedder):
31	"""
32	Thin wrapper over OpenAI embeddings.
33	Uses text-embedding-3-small by default (fast, 1536 dims).
34	"""
35
36	def __init__(
37	self,
38	api_key: Optional[str] = None,
39	model: str = "text-embedding-3-small",
40	client: Optional[OpenAI] = None,
41	):
42	# `client` lets callers (e.g. the Azure branch of create_embedder())
43	# inject a pre-built AzureOpenAI client — it's duck-type compatible
44	# since only .embeddings.create() is ever called on it.
45	self.client = client or OpenAI(api_key=api_key)
46	self.model = model
47
48	def _embed(self, inputs: List[str], max_retries: int = 5) -> List[List[float]]:
49	last_err = None
50	for attempt in range(max_retries):
51	try:
52	res = self.client.embeddings.create(model=self.model, input=inputs)
53	return [item.embedding for item in res.data]
54	except InternalServerError as e:
55	last_err = e
56	if attempt < max_retries - 1:
57	delay = 2 ** attempt
58	time.sleep(delay)
59	print(f"⚠️ OpenAI 500 error, retrying in {delay}s (attempt {attempt + 2}/{max_retries})...", flush=True)
60	else:
61	raise last_err from None
62	raise last_err from None
63
64	def embed_documents(self, texts: List[str]) -> List[List[float]]:
65	if not texts:
66	return []
67	return self._embed(texts)
68
69	def embed_query(self, text: str) -> List[float]:
70	return self._embed([text])[0]
71		26
72		27
73	# ── Provider selection ────────────────────────────────────────────────────	28	# ── Provider selection ────────────────────────────────────────────────────
⋯ 1 unchanged line hidden
74	#	29	#
75	# EMBEDDING_PROVIDER (env var) picks the embedding implementation. Defaults	30	# EMBEDDING_PROVIDER (env var) picks the embedding implementation. Defaults
76	# to today's behavior — on-prem/local deployments don't need to set anything.	31	# to today's behavior — on-prem/local deployments don't need to set anything.
77	#	32	#
78	# openai (default) — OpenAI's cloud embeddings API.	33	# openai (default) — OpenAIEmbeddings.
79	# azure_openai — Azure OpenAI Service embeddings. Requires	34	# azure_openai — AzureOpenAIEmbeddings. Requires AZURE_OPENAI_ENDPOINT,
80	# AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_KEY,	35	# AZURE_OPENAI_API_KEY, AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME.
81	# AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME (see .env).	36
		37	DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small"
		38
82		39
83	def create_embedder(	40	def create_embedder(
84	provider: Optional[str] = None,	41	provider: Optional[str] = None,
85	*,	42	*,
86	api_key: Optional[str] = None,	43	api_key: Optional[str] = None,
87	model: Optional[str] = None,	44	model: Optional[str] = None,
88	) -> Embedder:	45	cache_dir: Optional[str] = None,
89	"""Factory for the embedding model, selected by EMBEDDING_PROVIDER or `provider`."""	46	) -> Embeddings:
		47	"""Factory for the embedding model, selected by EMBEDDING_PROVIDER or `provider`.
		48
		49	cache_dir: when given, wraps the model in a CacheBackedEmbeddings backed by
		50	a local file store in that directory. Off by default so tests and one-off
		51	scripts never write cache files unless asked to.
		52	"""
90	provider = (provider or os.getenv("EMBEDDING_PROVIDER", "openai")).strip().lower()	53	provider = (provider or os.getenv("EMBEDDING_PROVIDER", "openai")).strip().lower()
91		54
92	if provider == "openai":	55	if provider == "openai":
		56	from langchain_openai import OpenAIEmbeddings
		57
93	resolved_key = api_key or os.getenv("OPENAI_API_KEY")	58	resolved_key = api_key or os.getenv("OPENAI_API_KEY")
94	if not resolved_key:	59	if not resolved_key:
95	raise RuntimeError("OPENAI_API_KEY must be set for EMBEDDING_PROVIDER=openai")	60	raise RuntimeError("OPENAI_API_KEY must be set for EMBEDDING_PROVIDER=openai")
96	return OpenAIEmbedder(api_key=resolved_key, model=model or "text-embedding-3-small")	61	model_name = model or DEFAULT_EMBEDDING_MODEL
		62	embedder: Embeddings = OpenAIEmbeddings(model=model_name, api_key=resolved_key, max_retries=5)
97		63
98	if provider == "azure_openai":	64	elif provider == "azure_openai":
99	from .model import build_azure_openai_client	65	from langchain_openai import AzureOpenAIEmbeddings
100		66
101	deployment = model or os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME")	67	deployment = model or os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME")
102	if not deployment:	68	if not deployment:
⋯ 2 unchanged lines hidden
103	raise RuntimeError(	69	raise RuntimeError(
104	"EMBEDDING_PROVIDER=azure_openai requires "	70	"EMBEDDING_PROVIDER=azure_openai requires "
105	"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME (the deployment name you "	71	"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME (the deployment name you "
106	"gave the embedding model in Azure — not the underlying model name)."	72	"gave the embedding model in Azure — not the underlying model name)."
107	)	73	)
108	return OpenAIEmbedder(client=build_azure_openai_client(api_key), model=deployment)	74	endpoint, resolved_key, api_version = azure_openai_settings(api_key)
		75	model_name = deployment
		76	embedder = AzureOpenAIEmbeddings(
		77	azure_endpoint=endpoint,
		78	api_key=resolved_key,
		79	api_version=api_version,
		80	azure_deployment=deployment,
		81	max_retries=5,
		82	)
109		83
110	raise ValueError(	84	else:
111	f"Unknown EMBEDDING_PROVIDER={provider!r}. Supported: openai, azure_openai."	85	raise ValueError(
		86	f"Unknown EMBEDDING_PROVIDER={provider!r}. Supported: openai, azure_openai."
		87	)
		88
		89	if cache_dir:
		90	return with_embedding_cache(embedder, cache_dir=cache_dir, namespace=f"{provider}:{model_name}")
		91	return embedder
		92
		93
		94	def with_embedding_cache(embedder: Embeddings, *, cache_dir: str, namespace: str) -> Embeddings:
		95	"""Wrap an embedder so repeated texts are served from a local file cache.
		96
		97	The namespace includes provider + model: vectors from different models are
		98	not interchangeable, so they must never share cache entries.
		99	query_embedding_cache=True caches query embeddings too (same store), so a
		100	repeated question skips the embedding call as well.
		101	"""
		102	from langchain_classic.embeddings import CacheBackedEmbeddings
		103	from langchain_classic.storage import LocalFileStore
		104
		105	# The namespace becomes part of every cache key, and LocalFileStore only
		106	# accepts [A-Za-z0-9_.-/] in keys — "openai:text-embedding-3-small" would
		107	# be rejected on the first write. Replace anything else (and "/", which
		108	# would create subdirectories) with "_".
		109	safe_namespace = re.sub(r"[^A-Za-z0-9_.\-]", "_", namespace)
		110
		111	store = LocalFileStore(cache_dir)
		112	return CacheBackedEmbeddings.from_bytes_store(
		113	embedder,
		114	store,
		115	namespace=safe_namespace,
		116	query_embedding_cache=True,
		117	# Keys are a hash of the text. The default (SHA-1) isn't
		118	# collision-resistant and LangChain warns about it; SHA-256 is.
		119	key_encoder="sha256",
112	)	120	)

03 agentic_chatbot.py

src/business/chatbot/agentic_chatbot.py · 306 → 285 lines · 273 rows changed

The agent loop: call the model, run any tools it asks for, repeat until it answers. Same loop, but LangGraph turns it into a graph.

cortex-core · hand-written

A while True: loop in _agent_loop(): call chat.completions.create(tools=...), run each tool call, append the results, repeat.
No cap on the number of rounds, so a model that keeps asking for tools loops forever.
Tools are a hand-written JSON schema list (TOOLS) plus an if tool_name == ... dispatcher.
Returns the answer text only; token usage is never collected.

cortex-langgraph · LangChain / LangGraph

A StateGraph with two nodes: agent (the model with tools bound) and tools (LangGraph's prebuilt ToolNode).
tools_condition routes to tools or END; recursion_limit caps it at 10 tool rounds.
MessagesState appends each step's messages, so nodes only return what's new.
Sums usage_metadata across every model call in the turn into last_usage.

Why it matters: The loop is identical in behavior, but the graph is bounded, can be drawn, streamed step by step and traced. Memory (Redis, Chroma, SQLite) is unchanged.

StateGraphMessagesStateToolNodetools_conditionbind_tools

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/chatbot/agentic_chatbot.py		cortex-langgraph src/business/chatbot/agentic_chatbot.py
1	# This cahtbot have both vectorDB and SQL database
2	# SQL Database -> for structured queries (e.g., task management, user profiles)
3	# VectorDB -> for free-form queries (e.g., knowledge retrieval, document search)
4	# This chatbot should use the SQL database of each user like last location, last search, etc.
5	# to make the conversation more personalized and context-aware. For example, when user ask about the trip suggestion
6	# it can use the last location of the user to suggest nearby attractions or restaurants.
7	"""	1	"""
8	Agentic chatbot — V3.	2	Agentic chatbot — LangGraph edition.
9		3
10	Memory architecture:	4	Memory architecture (unchanged from cortex-core):
11	Short-term → RedisMemory (recent turns, auto-expires)	5	Short-term → RedisMemory (recent turns, auto-expires)
12	Long-term → LongTermMemory / VectorDB (semantic search over all history)	6	Long-term → LongTermMemory (semantic search over all history, Chroma)
		7	Persistent → ChatHistoryManager (SQLite, survives restarts; feeds the history API)
13		8
14	Tool swap vs V2:	9	What changed: the tool-calling loop.
15	Removed: search_chat_history (SQL LIKE keyword search)
16	Added: search_vector_db (ChromaDB semantic / embedding search)
17		10
18	After every chat() turn the exchange is persisted to the vector store so	11	cortex-core ran a hand-written `while True:` loop — call the model, if it asked
19	future turns can recall it semantically.	12	for tools run them and append the results, repeat until it answers in plain
		13	text. Here that loop is a LangGraph `StateGraph` with two nodes:
		14
		15	┌─────────┐ tools_condition: model asked for tools? ┌─────────┐
		16	START │ agent │ ───────────────────── yes ────────────────▶ │ tools │
		17	────▶ │ (LLM) │ ◀────────────────────────────────────────── │ToolNode │
		18	└────┬────┘ └─────────┘
		19	│ no (plain answer)
		20	▼
		21	END
		22
		23	- `agent` calls the chat model with the tools bound (`llm.bind_tools`).
		24	- `tools` is LangGraph's prebuilt `ToolNode`: it runs every tool call in the
		25	last AI message and appends one ToolMessage per call.
		26	- `tools_condition` is the prebuilt router that picks between the two.
		27
		28	The loop is the same; what the graph adds is that it's inspectable (you can
		29	draw it, stream each step, trace it) and bounded — LangGraph stops a runaway
		30	loop at `recursion_limit` (see MAX_TOOL_ROUNDS) instead of spinning forever.
		31
		32	Why short-term memory is NOT a LangGraph checkpointer: the official Redis
		33	checkpointer needs Redis with the JSON and Search modules (Redis Stack / Redis 8
		34	with the query engine). Plain Redis — what Docker Compose, Homebrew and Azure
		35	Cache for Redis Basic/Standard provide — doesn't have them. So conversation
		36	turns stay in plain Redis lists, and are converted to LangChain messages when
		37	they enter the graph.
20	"""	38	"""
21		39
22	from __future__ import annotations	40	from __future__ import annotations
23		41
24	import json
25	import os
26	from typing import Dict, List, Optional	42	from typing import Dict, List, Optional
27		43
28	from dotenv import load_dotenv	44	from dotenv import load_dotenv
29	from openai import OpenAI	45	from langchain_core.language_models import BaseChatModel
		46	from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, convert_to_messages
		47	from langgraph.graph import START, MessagesState, StateGraph
		48	from langgraph.prebuilt import ToolNode, tools_condition
30		49
31	from src.business.core.model import build_azure_openai_client	50	from src.business.chatbot.tools import build_tools
		51	from src.business.core.live_data import LiveDataProvider, create_live_data_provider
		52	from src.business.core.model import create_llm, model_label, resolve_provider
32	from src.business.core.prompt_builder import build_agentic_system_prompt	53	from src.business.core.prompt_builder import build_agentic_system_prompt
33	from src.business.core.live_data import create_live_data_provider, LiveDataProvider
34	from src.memory.chat_history_manager import ChatHistoryManager	54	from src.memory.chat_history_manager import ChatHistoryManager
35	from src.memory.long_term_memory import LongTermMemory	55	from src.memory.long_term_memory import LongTermMemory
36	from src.memory.redis_memory import RedisMemory	56	from src.memory.redis_memory import ShortTermMemoryBase
37	from src.utils.config import load_config	57	from src.utils.config import load_config
		58
		59	# Providers whose chat models support OpenAI-style tool calling.
		60	TOOL_CALLING_PROVIDERS = ("openai", "azure_openai")
		61
		62	# Upper bound on agent → tools round-trips per user turn. cortex-core's
		63	# while-loop had no cap at all, and LangGraph's own default (10,007 steps) is
		64	# effectively none either. Each round is two graph steps (agent, then tools);
		65	# the final answer needs two more (the answering agent step, and the step that
		66	# reaches END). Exceeding it raises GraphRecursionError.
		67	# It means LangGraph puts a safety limit on how many times the agent can call tools in one user request.
		68	#If you configure max_tool_rounds = 5, the agent can make at most 5 tool-call rounds.
		69	#If it exceeds the limit, LangGraph raises GraphRecursionError.
		70
		71	MAX_TOOL_ROUNDS = 10
		72	RECURSION_LIMIT = 2 * MAX_TOOL_ROUNDS + 2
38		73
39		74
40	class AgenticChatbot:	75	class AgenticChatbot:
⋯ 3 unchanged lines hidden
41	"""	76	"""
42	Conversational agent with semantic long-term memory.	77	Conversational agent with semantic long-term memory.
43		78
44	Dependencies are injected so the chatbot is storage-agnostic and testable.	79	Dependencies are injected so the chatbot is storage-agnostic and testable.
45		80
46	Args:	81	Args:
47	long_term_memory: LongTermMemory instance (wraps VectorDB + embedder).	82	long_term_memory: LongTermMemory instance (wraps a LangChain vector store).
48	redis_memory: RedisMemory instance for short-term session context.	83	redis_memory: Short-term session memory (RedisMemory in production).
49	user_id: Stable identifier for the current user.	84	user_id: Stable identifier for the current user.
		85	chat_history_manager: Optional SQLite history (persistent, feeds the API).
50	user_info: Optional dict of user profile data injected into the	86	user_info: Optional dict of user profile data injected into the
51	system prompt (name, location, preferences, …).	87	system prompt (name, location, preferences, …).
52	api_key: OpenAI API key. Falls back to OPENAI_API_KEY env var.	88	api_key: Provider API key. Falls back to the provider's env var.
53	model_name: Chat model. Falls back to config → "gpt-4o".	89	model_name: Chat model / Azure deployment. Falls back to config → "gpt-4o".
		90	live_data_provider: Backend for the web_search tool.
		91	llm: A ready chat model. Injected by tests; when omitted the
		92	model is built from LLM_PROVIDER via create_llm().
54	"""	93	"""
55		94
56	# ------------------------------------------------------------------ tools
57	TOOLS: List[Dict] = [
58	{
59	"type": "function",
60	"function": {
61	"name": "search_vector_db",
62	"description": (
63	"Semantic search over all past conversations. "
64	"Use this when you need to recall what the user said, "
65	"asked, or was told in a previous session."
66	),
67	"parameters": {
68	"type": "object",
69	"properties": {
70	"query": {
71	"type": "string",
72	"description": "Natural-language description of what to look for.",
73	}
74	},
75	"required": ["query"],
76	},
77	},
78	},
79	{
80	"type": "function",
81	"function": {
82	"name": "web_search",
83	"description": (
84	"Search the web for current information, news, or real-time data. "
85	"Use this when the user asks about current events, recent news, "
86	"or information that changes frequently (e.g., 'What's in the news?', "
87	"'latest updates on X topic')."
88	),
89	"parameters": {
90	"type": "object",
91	"properties": {
92	"query": {
93	"type": "string",
94	"description": "Search query (e.g., 'latest AI news', 'weather NYC', 'trending today')",
95	}
96	},
97	"required": ["query"],
98	},
99	},
100	}
101	]
102
103	# ------------------------------------------------------------ constructor
104	def __init__(	95	def __init__(
105	self,	96	self,
106	long_term_memory: LongTermMemory,	97	long_term_memory: LongTermMemory,
107	redis_memory: RedisMemory,	98	redis_memory: ShortTermMemoryBase,
108	user_id: str,	99	user_id: str,
109	chat_history_manager: Optional["ChatHistoryManager"] = None,	100	chat_history_manager: Optional[ChatHistoryManager] = None,
110	user_info: Optional[Dict] = None,	101	user_info: Optional[Dict] = None,
111	api_key: Optional[str] = None,	102	api_key: Optional[str] = None,
112	model_name: Optional[str] = None,	103	model_name: Optional[str] = None,
113	live_data_provider: Optional[LiveDataProvider] = None,	104	live_data_provider: Optional[LiveDataProvider] = None,
		105	llm: Optional[BaseChatModel] = None,
114	) -> None:	106	) -> None:
115	load_dotenv()	107	load_dotenv()
116		108
⋯ 3 unchanged lines hidden
117	self.long_term_memory = long_term_memory	109	self.long_term_memory = long_term_memory
118	self.redis_memory = redis_memory	110	self.redis_memory = redis_memory
119	self.user_id = user_id	111	self.user_id = user_id
120	self.chat_history_manager = chat_history_manager	112	self.chat_history_manager = chat_history_manager
121	self.user_info = user_info or {}	113	self.user_info = user_info or {}
122	self.live_data_provider = live_data_provider or create_live_data_provider()	114	self.live_data_provider = live_data_provider or create_live_data_provider()
		115	self.last_usage: Dict[str, int] = {}
123		116
124	config = load_config()	117	if llm is None:
125	llm_config = config.get("llm_config") or {}	118	llm = self._build_llm(api_key=api_key, model_name=model_name)
		119	self.llm = llm
		120	self.model_name = model_label(llm)
126		121
127	# LLM_PROVIDER selects the client for this tool-calling loop, same	122	self.tools = build_tools(
128	# switch used by create_llm() elsewhere. Local Hugging Face models	123	long_term_memory=self.long_term_memory,
129	# don't support this agent's function-calling flow, so that provider	124	live_data_provider=self.live_data_provider,
130	# isn't valid here.	125	user_id=self.user_id,
131	provider = os.getenv("LLM_PROVIDER", "openai").strip().lower()	126	)
		127	self.graph = self._build_graph()
132		128
133	if provider == "openai":	129	# ------------------------------------------------------------- the model
134	resolved_key = api_key or os.getenv("OPENAI_API_KEY")	130	@staticmethod
135	if not resolved_key:	131	def _build_llm(api_key: Optional[str], model_name: Optional[str]) -> BaseChatModel:
136	raise RuntimeError("OPENAI_API_KEY must be set in environment or passed to __init__")	132	# LLM_PROVIDER selects the model, the same switch used by create_llm()
137	self.client = OpenAI(api_key=resolved_key)	133	# everywhere else. Local Hugging Face models don't support this agent's
138	self.model_name = (	134	# tool-calling flow, so that provider isn't valid here — fail loudly
139	model_name	135	# rather than silently answering without tools.
140	or llm_config.get("chat_model")	136	provider = resolve_provider()
141	or "gpt-4o"	137	if provider not in TOOL_CALLING_PROVIDERS:
142	)
143
144	elif provider == "azure_openai":
145	deployment = model_name or os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME")
146	if not deployment:
147	raise RuntimeError(
148	"LLM_PROVIDER=azure_openai requires AZURE_OPENAI_CHAT_DEPLOYMENT_NAME."
149	)
150	self.client = build_azure_openai_client(api_key)
151	self.model_name = deployment
152
153	else:
154	raise NotImplementedError(	138	raise NotImplementedError(
155	f"AgenticChatbot's tool-calling loop only supports LLM_PROVIDER in "	139	f"AgenticChatbot's tool-calling graph only supports LLM_PROVIDER in "
156	f"('openai', 'azure_openai') — got {provider!r}. Local Hugging Face "	140	f"{TOOL_CALLING_PROVIDERS} — got {provider!r}. Local Hugging Face "
157	"models don't support this agent's function-calling flow."	141	"models don't support this agent's function-calling flow."
158	)	142	)
159		143
160	# ---------------------------------------------------------- tool dispatch	144	if provider == "openai":
161	def _handle_tool_call(self, tool_name: str, tool_args: Dict) -> str:	145	config = load_config()
162	if tool_name == "search_vector_db":	146	llm_config = config.get("llm_config") or {}
163	results = self.long_term_memory.recall(	147	model_name = model_name or llm_config.get("chat_model") or "gpt-4o"
164	query=tool_args["query"],	148	# For azure_openai, model_name is the deployment; create_llm falls back
165	user_id=self.user_id,	149	# to AZURE_OPENAI_CHAT_DEPLOYMENT_NAME when it's None.
166	)	150	return create_llm(provider, model_name=model_name, api_key=api_key)
167	if not results:
168	return "No relevant past conversations found."
169	return "\n\n".join(r["text"] for r in results)
170		151
171	if tool_name == "web_search":	152	# ------------------------------------------------------------- the graph
172	try:	153	def _build_graph(self):
173	results = self.live_data_provider.search(	154	"""Compile the agent ⇄ tools graph (replaces cortex-core's while-loop)."""
174	query=tool_args["query"],	155	llm_with_tools = self.llm.bind_tools(self.tools)
175	limit=5,
176	)
177	if not results:
178	return f"No search results found for '{tool_args['query']}'."
179		156
180	# Format results as readable text	157	def agent(state: MessagesState) -> Dict[str, List[BaseMessage]]:
181	formatted_results = []	158	# MessagesState's reducer APPENDS returned messages to the state,
182	for i, result in enumerate(results, 1):	159	# so each step only returns what's new.
183	title = result.get("title", "")	160	return {"messages": [llm_with_tools.invoke(state["messages"])]}
184	summary = result.get("summary", "")
185	source = result.get("source", "Unknown")
186	url = result.get("url", "")
187		161
188	text = f"{i}. {title}\n Source: {source}\n {summary}"	162	builder = StateGraph(MessagesState)
189	if url:	163	builder.add_node("agent", agent)
190	text += f"\n URL: {url}"	164	builder.add_node("tools", ToolNode(self.tools))
191	formatted_results.append(text)	165	builder.add_edge(START, "agent")
192		166	builder.add_conditional_edges("agent", tools_condition) # → "tools" or END
193	return "\n\n".join(formatted_results)	167	builder.add_edge("tools", "agent")
194	except Exception as e:	168	return builder.compile().with_config(recursion_limit=RECURSION_LIMIT)
195	return f"Web search failed: {str(e)}"
196
197	return f"Unknown tool: {tool_name}"
198
199	# ------------------------------------------------------------ agent loop
200	def _agent_loop(self, messages: List[Dict]) -> str:
201	"""
202	Runs the ReAct loop: call LLM → handle tool calls → repeat until the
203	model returns a plain text response.
204	"""
205	while True:
206	response = self.client.chat.completions.create(
207	model=self.model_name,
208	messages=messages,
209	tools=self.TOOLS,
210	tool_choice="auto",
211	)
212	msg = response.choices[0].message
213
214	if msg.tool_calls:
215	# Append assistant message with tool_calls intact
216	messages.append(msg)
217	for tc in msg.tool_calls:
218	result = self._handle_tool_call(
219	tc.function.name,
220	json.loads(tc.function.arguments),
221	)
222	messages.append({
223	"role": "tool",
224	"tool_call_id": tc.id,
225	"content": result,
226	})
227	else:
228	return msg.content
229		169
230	# --------------------------------------------------------------- chat API	170	# --------------------------------------------------------------- chat API
231	async def chat(	171	async def chat(
⋯ 8 unchanged lines hidden
232	self,	172	self,
233	user_message: str,	173	user_message: str,
234	session_id: Optional[str] = None,	174	session_id: Optional[str] = None,
235	chat_summary: Optional[str] = None,	175	chat_summary: Optional[str] = None,
236	) -> str:	176	) -> str:
237	"""	177	"""
238	Process one user turn end-to-end.	178	Process one user turn end-to-end.
239		179
240	1. Fetch recent turns from Redis (short-term window).	180	1. Fetch recent turns from Redis (short-term window).
241	2. Semantic pre-fetch from vector store for likely-relevant context.	181	2. Semantic pre-fetch from vector store for likely-relevant context.
242	3. Build system prompt (user_info + summary + vector snippets).	182	3. Build system prompt (user_info + summary + vector snippets).
243	4. Run agent loop (may call search_vector_db on demand).	183	4. Run the agent graph (may call search_vector_db / web_search on demand).
244	5. Persist the exchange to Redis and to the vector store.	184	5. Persist the exchange to Redis, the vector store, and SQLite.
245		185
246	Args:	186	Args:
247	user_message: The user's raw input.	187	user_message: The user's raw input.
⋯ 1 unchanged line hidden
248	session_id: Redis key prefix for this session (defaults to user_id).	188	session_id: Redis key prefix for this session (defaults to user_id).
249	chat_summary: Optional rolling summary of the conversation so far.	189	chat_summary: Optional rolling summary of the conversation so far.
250		190
251	Returns:	191	Returns:
252	The assistant's final response string.	192	The assistant's final response string. Token usage for the turn
		193	(summed over every model call the graph made) is left on
		194	`self.last_usage`.
253	"""	195	"""
254	session_id = session_id or self.user_id	196	session_id = session_id or self.user_id
255		197
⋯ 20 unchanged lines hidden
256	# 1. Short-term context from Redis (fast, in-session window)	198	# 1. Short-term context from Redis (fast, in-session window)
257	recent_turns = await self.redis_memory.get_messages(session_id)	199	recent_turns = await self.redis_memory.get_messages(session_id)
258		200
259	# 1b. If this is a fresh Redis session (no turns yet) and we have a	201	# 1b. If this is a fresh Redis session (no turns yet) and we have a
260	# SQLite manager, pull the previous session's messages as a	202	# SQLite manager, pull the previous session's messages as a
261	# rolling summary so the model has prior context on restart.	203	# rolling summary so the model has prior context on restart.
262	if not recent_turns and chat_summary is None and self.chat_history_manager:	204	if not recent_turns and chat_summary is None and self.chat_history_manager:
263	chat_summary = self.chat_history_manager.get_latest_summary(self.user_id)	205	chat_summary = self.chat_history_manager.get_latest_summary(self.user_id)
264		206
265	# 2. Proactive semantic prefetch — surface relevant long-term memories	207	# 2. Proactive semantic prefetch — surface relevant long-term memories
266	# before the agent even has a chance to call the tool.	208	# before the agent even has a chance to call the tool.
267	vector_results = self.long_term_memory.recall(	209	vector_results = self.long_term_memory.recall(
268	query=user_message,	210	query=user_message,
269	user_id=self.user_id,	211	user_id=self.user_id,
270	)	212	)
271		213
272	# 3. Build the system prompt	214	# 3. Build the system prompt
273	system_prompt = build_agentic_system_prompt(	215	system_prompt = build_agentic_system_prompt(
274	user_info=self.user_info,	216	user_info=self.user_info,
275	vector_results=vector_results,	217	vector_results=vector_results,
276	chat_summary=chat_summary,	218	chat_summary=chat_summary,
277	)	219	)
278		220
279	# 4. Assemble message list and run the agent loop	221	# 4. Run the graph. Redis stores plain {"role", "content"} dicts;
280	messages: List[Dict] = [{"role": "system", "content": system_prompt}]	222	# convert_to_messages turns them into Human/AI messages.
281	messages.extend(recent_turns)	223	messages: List[BaseMessage] = [SystemMessage(content=system_prompt)]
282	messages.append({"role": "user", "content": user_message})	224	messages.extend(convert_to_messages(recent_turns))
		225	messages.append(HumanMessage(content=user_message))
283		226
284	response = self._agent_loop(messages)	227	result = await self.graph.ainvoke({"messages": messages})
		228	new_messages = result["messages"][len(messages):]
		229	response = self._final_text(new_messages)
		230	self.last_usage = self._sum_usage(new_messages)
285		231
286	# 5a. Persist to short-term (Redis) — fast, auto-expires	232	# 5a. Persist to short-term (Redis) — fast, auto-expires
287	await self.redis_memory.add_message(session_id, "user", user_message)	233	await self.redis_memory.add_message(session_id, "user", user_message)
⋯ 16 unchanged lines hidden
288	await self.redis_memory.add_message(session_id, "assistant", response)	234	await self.redis_memory.add_message(session_id, "assistant", response)
289		235
290	# 5b. Persist to long-term (vector store) — makes this turn searchable	236	# 5b. Persist to long-term (vector store) — makes this turn searchable
291	self.long_term_memory.remember_conversation(	237	self.long_term_memory.remember_conversation(
292	user_message=user_message,	238	user_message=user_message,
293	assistant_response=response,	239	assistant_response=response,
294	user_id=self.user_id,	240	user_id=self.user_id,
295	)	241	)
296		242
297	# 5c. Persist to SQLite — survives Redis TTL expiry and restarts	243	# 5c. Persist to SQLite — survives Redis TTL expiry and restarts
298	if self.chat_history_manager:	244	if self.chat_history_manager:
299	# Use the first user message as the session title	245	# Use the first user message as the session title
300	is_first_turn = not recent_turns	246	is_first_turn = not recent_turns
301	title = user_message[:60] if is_first_turn else "New conversation"	247	title = user_message[:60] if is_first_turn else "New conversation"
302	self.chat_history_manager.ensure_session(session_id, self.user_id, title)	248	self.chat_history_manager.ensure_session(session_id, self.user_id, title)
303	self.chat_history_manager.save_message(session_id, self.user_id, "user", user_message)	249	self.chat_history_manager.save_message(session_id, self.user_id, "user", user_message)
304	self.chat_history_manager.save_message(session_id, self.user_id, "assistant", response)	250	self.chat_history_manager.save_message(session_id, self.user_id, "assistant", response)
305		251
306	return response	252	return response
		253
		254	# ---------------------------------------------------------------- helpers
		255	#These are two helper functions that process the messages produced by the LangGraph run.
		256	# extract the final answer from the model.
		257	@staticmethod
		258	def _final_text(new_messages: List[BaseMessage]) -> str:
		259	"""The graph ends on the model's plain-text answer — the last AIMessage."""
		260	for message in reversed(new_messages):
		261	if isinstance(message, AIMessage):
		262	# .text flattens both plain-string and content-block responses.
		263	return message.text
		264	return ""
		265
		266	@staticmethod
		267	#Purpose: calculate the total token usage for the entire user turn.
		268	def _sum_usage(new_messages: List[BaseMessage]) -> Dict[str, int]:
		269	"""Add up token usage across every model call made during this turn.
		270
		271	A turn with two tool rounds calls the model three times; cortex-core
		272	never reported tokens at all, so cost tracking could never fire.
		273	"""
		274	totals = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
		275	for message in new_messages:
		276	usage = getattr(message, "usage_metadata", None)
		277	if usage:
		278	for key in totals:
		279	totals[key] += int(usage.get(key, 0) or 0)
		280	return totals
		281
		282
		283	# So these functions are not LangGraph orchestration.
		284	# They are application-level helpers used after/during
		285	# the graph execution to extract the result and track usage/cost.

04 tools.py

src/business/chatbot/tools.py · 85 → 75 lines · 82 rows changed

The agent's two tools, search_vector_db and web_search. New file in cortex-langgraph; the left side is the cortex-core code it replaces.

cortex-core · hand-written

Each tool existed twice: once as a JSON schema in AgenticChatbot.TOOLS, once as a branch in _handle_tool_call().
Keeping schema, name and arguments in sync between the two was manual.

cortex-langgraph · LangChain / LangGraph

Each tool is a plain function with @tool. LangChain builds the JSON schema from the signature and docstring.
LangGraph's ToolNode does the dispatch, so there's no if/elif chain.
build_tools() gives the tools their user, memory store and search provider by closure.

Why it matters: The function is the tool, so the schema can't drift from the code. Tool names and descriptions are unchanged, so the model sees exactly the same tools.

@toolBaseToolToolNode

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/chatbot/agentic_chatbot.py (excerpt: TOOLS + _handle_tool_call)		cortex-langgraph src/business/chatbot/tools.py
56	# ------------------------------------------------------------------ tools	1	"""The agent's tools — LangChain `@tool` edition.
57	TOOLS: List[Dict] = [	2
58	{	3	In cortex-core each tool existed twice: once as a hand-written JSON schema in
59	"type": "function",	4	`AgenticChatbot.TOOLS`, and once as a branch of `_handle_tool_call()`'s
60	"function": {	5	`if tool_name == ...` dispatch. Keeping the two in sync was manual.
61	"name": "search_vector_db",	6
62	"description": (	7	With `@tool`, the function IS the tool: LangChain derives the JSON schema from
63	"Semantic search over all past conversations. "	8	the signature + docstring, and LangGraph's `ToolNode` does the dispatch. The
64	"Use this when you need to recall what the user said, "	9	tool names and descriptions below are unchanged from cortex-core, so the model
65	"asked, or was told in a previous session."	10	sees exactly the same tool definitions.
66	),	11
67	"parameters": {	12	The tools need per-request context (which user, which memory store, which
68	"type": "object",	13	search provider). They get it by closure: `build_tools()` is called once per
69	"properties": {	14	chatbot, and the inner functions capture those objects.
70	"query": {	15	"""
71	"type": "string",	16
72	"description": "Natural-language description of what to look for.",	17	from __future__ import annotations
73	}	18
74	},	19	from typing import List
75	"required": ["query"],	20
76	},	21	from langchain_core.tools import BaseTool, tool
77	},	22
78	},	23	from src.business.core.live_data import LiveDataProvider
79	{	24	from src.memory.long_term_memory import LongTermMemory
80	"type": "function",	25
81	"function": {	26
82	"name": "web_search",	27	def build_tools(
83	"description": (	28	long_term_memory: LongTermMemory,
84	"Search the web for current information, news, or real-time data. "	29	live_data_provider: LiveDataProvider,
85	"Use this when the user asks about current events, recent news, "	30	user_id: str,
86	"or information that changes frequently (e.g., 'What's in the news?', "	31	) -> List[BaseTool]:
87	"'latest updates on X topic')."	32	"""Create the agent's tools, bound to one user's memory and a search provider."""
88	),	33
89	"parameters": {	34	@tool
90	"type": "object",	35	def search_vector_db(query: str) -> str:
91	"properties": {	36	"""Semantic search over all past conversations. Use this when you need to recall what the user said, asked, or was told in a previous session.
92	"query": {	37
93	"type": "string",	38	Args:
94	"description": "Search query (e.g., 'latest AI news', 'weather NYC', 'trending today')",	39	query: Natural-language description of what to look for.
95	}	40	"""
96	},	41	results = long_term_memory.recall(query=query, user_id=user_id)
97	"required": ["query"],	42	if not results:
98	},	43	return "No relevant past conversations found."
99	},	44	return "\n\n".join(r["text"] for r in results)
100	}	45
101	]	46	@tool
161	def _handle_tool_call(self, tool_name: str, tool_args: Dict) -> str:	47	def web_search(query: str) -> str:
162	if tool_name == "search_vector_db":	48	"""Search the web for current information, news, or real-time data. Use this when the user asks about current events, recent news, or information that changes frequently (e.g., 'What's in the news?', 'latest updates on X topic').
163	results = self.long_term_memory.recall(	49
164	query=tool_args["query"],	50	Args:
165	user_id=self.user_id,	51	query: Search query (e.g., 'latest AI news', 'weather NYC', 'trending today')
166	)	52	"""
		53	try:
		54	results = live_data_provider.search(query=query, limit=5)
167	if not results:	55	if not results:
168	return "No relevant past conversations found."	56	return f"No search results found for '{query}'."
169	return "\n\n".join(r["text"] for r in results)
170		57
171	if tool_name == "web_search":	58	# Format results as readable text
172	try:	59	formatted_results = []
173	results = self.live_data_provider.search(	60	for i, result in enumerate(results, 1):
174	query=tool_args["query"],	61	title = result.get("title", "")
175	limit=5,	62	summary = result.get("summary", "")
176	)	63	source = result.get("source", "Unknown")
177	if not results:	64	url = result.get("url", "")
178	return f"No search results found for '{tool_args['query']}'."
179		65
180	# Format results as readable text	66	text = f"{i}. {title}\n Source: {source}\n {summary}"
181	formatted_results = []	67	if url:
182	for i, result in enumerate(results, 1):	68	text += f"\n URL: {url}"
183	title = result.get("title", "")	69	formatted_results.append(text)
184	summary = result.get("summary", "")
185	source = result.get("source", "Unknown")
186	url = result.get("url", "")
187		70
188	text = f"{i}. {title}\n Source: {source}\n {summary}"	71	return "\n\n".join(formatted_results)
189	if url:	72	except Exception as e:
190	text += f"\n URL: {url}"	73	return f"Web search failed: {str(e)}"
191	formatted_results.append(text)
192		74
193	return "\n\n".join(formatted_results)	75	return [search_vector_db, web_search]
194	except Exception as e:
195	return f"Web search failed: {str(e)}"
196
197	return f"Unknown tool: {tool_name}"
198
199	# ------------------------------------------------------------ agent loop

05 chatbot/__init__.py

src/business/chatbot/__init__.py · 129 → 124 lines · 30 rows changed

Wiring: builds the chatbot and its memory for each request. Small diff, but it's where the real token counts reach the API.

cortex-core · hand-written

Builds the vector store, then passes the embedder and an _IdentityChunker to LongTermMemory.
The API response hardcodes model_used: "gpt-4o" and tokens_used: None.

cortex-langgraph · LangChain / LangGraph

The embedder goes into the vector store instead, and LongTermMemory(vectorstore=...) needs nothing else.
Returns the real model name plus tokens_used, input_tokens and output_tokens from chatbot.last_usage.

Why it matters: With real token counts, the controller's cost-tracking branch finally runs.

dependency wiringusage_metadata

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/chatbot/__init__.py		cortex-langgraph src/business/chatbot/__init__.py
⋯ 8 unchanged lines hidden
1	"""	1	"""
2	Business logic entry-points for the chat API.	2	Business logic entry-points for the chat API.
3		3
4	process_chat_message → called by ChatController.send_message	4	process_chat_message → called by ChatController.send_message
5	get_chat_history → called by ChatController.get_chat_history	5	get_chat_history → called by ChatController.get_chat_history
6	"""	6	"""
7		7
8	from __future__ import annotations	8	from __future__ import annotations
9		9
10	import os	10	import os
11	from pathlib import Path	11	from pathlib import Path
12	from typing import Dict, List	12	from typing import Dict
13		13
14	from dotenv import load_dotenv	14	from dotenv import load_dotenv
15		15
⋯ 19 unchanged lines hidden
16	from src.business.chatbot.agentic_chatbot import AgenticChatbot	16	from src.business.chatbot.agentic_chatbot import AgenticChatbot
17	from src.database.dto import ChatMessageRequest, ChatHistoryRequest	17	from src.database.dto import ChatMessageRequest, ChatHistoryRequest
18	from src.memory.chat_history_manager import ChatHistoryManager	18	from src.memory.chat_history_manager import ChatHistoryManager
19	from src.memory.long_term_memory import LongTermMemory	19	from src.memory.long_term_memory import LongTermMemory
20	from src.memory.redis_memory import create_memory	20	from src.memory.redis_memory import create_memory
21	from src.memory.vectordb import create_conversation_vector_store	21	from src.memory.vectordb import create_conversation_vector_store
22	from src.business.core.embedding import create_embedder	22	from src.business.core.embedding import create_embedder
23	from src.utils.config import load_config	23	from src.utils.config import load_config
24		24
25	load_dotenv()	25	load_dotenv()
26		26
27	# Project root = src/business/chatbot/__init__.py → up 3 levels	27	# Project root = src/business/chatbot/__init__.py → up 3 levels
28	_PROJECT_ROOT = Path(__file__).resolve().parents[3]	28	_PROJECT_ROOT = Path(__file__).resolve().parents[3]
29		29
30		30
31	def _resolve(relative_path: str) -> Path:	31	def _resolve(relative_path: str) -> Path:
32	"""Resolve a config-relative path against the project root."""	32	"""Resolve a config-relative path against the project root."""
33	p = Path(relative_path)	33	p = Path(relative_path)
34	return p if p.is_absolute() else _PROJECT_ROOT / p	34	return p if p.is_absolute() else _PROJECT_ROOT / p
35		35
36		36
37	# ---------------------------------------------------------------------------	37	# ---------------------------------------------------------------------------
38	# Minimal chunker — conversation pairs are stored whole, so we never need
39	# to split them. Only used by LongTermMemory.remember(); the chatbot path
40	# always calls remember_conversation() which skips chunking.
41	# ---------------------------------------------------------------------------
42	class _IdentityChunker:
43	def split(self, text: str) -> List[str]:
44	return [text]
45
46
47	# ---------------------------------------------------------------------------
48	# Dependency factory	38	# Dependency factory
49	# ---------------------------------------------------------------------------	39	# ---------------------------------------------------------------------------
50	def _make_chatbot(user_id: str) -> AgenticChatbot:	40	def _make_chatbot(user_id: str) -> AgenticChatbot:
⋯ 10 unchanged lines hidden
51	cfg = load_config()	41	cfg = load_config()
52	dirs = cfg.get("directories", {})	42	dirs = cfg.get("directories", {})
53		43
54	api_key = os.getenv("OPENAI_API_KEY")	44	api_key = os.getenv("OPENAI_API_KEY")
55	redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")	45	redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
56		46
57	# Chroma persistent path — anchored to project root	47	# Chroma persistent path — anchored to project root
58	vectordb_dir = _resolve(os.getenv("VECTORDB_DIR") or dirs.get("vectordb_dir", "data/vectordb"))	48	vectordb_dir = _resolve(os.getenv("VECTORDB_DIR") or dirs.get("vectordb_dir", "data/vectordb"))
59	vectordb_dir.mkdir(parents=True, exist_ok=True)	49	vectordb_dir.mkdir(parents=True, exist_ok=True)
60		50
61	# SQLite path — anchored to project root	51	# SQLite path — anchored to project root
62	db_path = _resolve(os.getenv("SQLITE_DB_PATH") or dirs.get("db_path", "data/chatbot.db"))	52	db_path = _resolve(os.getenv("SQLITE_DB_PATH") or dirs.get("db_path", "data/chatbot.db"))
63		53
64	vectordb = create_conversation_vector_store(	54	# The vector store embeds text itself, so the embedder is handed to it
		55	# rather than to LongTermMemory (which is where cortex-core put it).
		56	embedder = create_embedder(api_key=api_key)
		57	vectorstore = create_conversation_vector_store(
65	collection_name="chat_history",	58	collection_name="chat_history",
		59	embedding=embedder,
66	persist_directory=str(vectordb_dir),	60	persist_directory=str(vectordb_dir),
67	)	61	)
68	embedder = create_embedder(api_key=api_key)	62	long_term_memory = LongTermMemory(vectorstore=vectorstore)
69	long_term_memory = LongTermMemory(
70	vectordb=vectordb,
71	embedder=embedder,
72	chunker=_IdentityChunker(),
73	)
74	redis_memory = create_memory(url=redis_url)	63	redis_memory = create_memory(url=redis_url)
75	chat_history_manager = ChatHistoryManager(db_path=str(db_path))	64	chat_history_manager = ChatHistoryManager(db_path=str(db_path))
76		65
⋯ 18 unchanged lines hidden
77	return AgenticChatbot(	66	return AgenticChatbot(
78	long_term_memory=long_term_memory,	67	long_term_memory=long_term_memory,
79	redis_memory=redis_memory,	68	redis_memory=redis_memory,
80	user_id=user_id,	69	user_id=user_id,
81	chat_history_manager=chat_history_manager,	70	chat_history_manager=chat_history_manager,
82	)	71	)
83		72
84		73
85	# ---------------------------------------------------------------------------	74	# ---------------------------------------------------------------------------
86	# Public API	75	# Public API
87	# ---------------------------------------------------------------------------	76	# ---------------------------------------------------------------------------
88	async def process_chat_message(request: ChatMessageRequest) -> Dict:	77	async def process_chat_message(request: ChatMessageRequest) -> Dict:
89	user_id = str(request.user_id) if request.user_id else "anonymous"	78	user_id = str(request.user_id) if request.user_id else "anonymous"
90	session_id = request.session_id or user_id	79	session_id = request.session_id or user_id
91		80
92	chatbot = _make_chatbot(user_id)	81	chatbot = _make_chatbot(user_id)
93	reply = await chatbot.chat(	82	reply = await chatbot.chat(
94	user_message=request.message,	83	user_message=request.message,
95	session_id=session_id,	84	session_id=session_id,
96	)	85	)
97		86
		87	# LangChain reports token usage on every model response, so — unlike
		88	# cortex-core, which always returned tokens_used=None — the controller now
		89	# receives real counts and its cost-tracking branch actually runs.
		90	usage = chatbot.last_usage
98	return {	91	return {
99	"reply": reply,	92	"reply": reply,
100	"model_used": "gpt-4o",	93	"model_used": chatbot.model_name,
101	"tokens_used": None,	94	"tokens_used": usage.get("total_tokens") or None,
		95	"input_tokens": usage.get("input_tokens", 0),
		96	"output_tokens": usage.get("output_tokens", 0),
102	}	97	}
103		98
104		99
⋯ 25 unchanged lines hidden
105	async def get_chat_history(request: ChatHistoryRequest) -> Dict:	100	async def get_chat_history(request: ChatHistoryRequest) -> Dict:
106	"""	101	"""
107	Returns messages for a session from SQLite — persistent across restarts.	102	Returns messages for a session from SQLite — persistent across restarts.
108	Falls back to an empty list gracefully if the session doesn't exist yet.	103	Falls back to an empty list gracefully if the session doesn't exist yet.
109	"""	104	"""
110	cfg = load_config()	105	cfg = load_config()
111	db_path = _resolve(	106	db_path = _resolve(
112	os.getenv("SQLITE_DB_PATH")	107	os.getenv("SQLITE_DB_PATH")
113	or cfg.get("directories", {}).get("db_path", "data/chatbot.db")	108	or cfg.get("directories", {}).get("db_path", "data/chatbot.db")
114	)	109	)
115	manager = ChatHistoryManager(db_path=str(db_path))	110	manager = ChatHistoryManager(db_path=str(db_path))
116		111
117	session_id = request.session_id or str(request.user_id)	112	session_id = request.session_id or str(request.user_id)
118	messages = manager.get_messages(	113	messages = manager.get_messages(
119	session_id=session_id,	114	session_id=session_id,
120	limit=request.limit,	115	limit=request.limit,
121	offset=request.offset,	116	offset=request.offset,
122	)	117	)
123	total = manager.count_messages(session_id)	118	total = manager.count_messages(session_id)
124		119
125	return {	120	return {
126	"messages": messages,	121	"messages": messages,
127	"total": total,	122	"total": total,
128	"session_id": session_id,	123	"session_id": session_id,
129	}	124	}

06 prompt_builder.py

src/business/core/prompt_builder.py · 141 → 152 lines · 156 rows changed

How prompts are assembled. Wording is identical on purpose so the two apps can be compared fairly; only the mechanism changes.

cortex-core · hand-written

A PromptBuilder class builds prompts with f-strings: build_prompt_text(), build_messages() (OpenAI message dicts) and build_prompt().

cortex-langgraph · LangChain / LangGraph

The same wording lives in LangChain templates: build_rag_prompt() returns a ChatPromptTemplate, and the agent prompt is a PromptTemplate.
format_context() formats LangChain Documents for the prompt.
A template is a Runnable, so the RAG graph can pipe it straight into the model: prompt | llm.

Why it matters: Templates compose with | and show up as their own step in LangSmith / Langfuse traces.

ChatPromptTemplatePromptTemplateLCEL

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/core/prompt_builder.py		cortex-langgraph src/business/core/prompt_builder.py
		1	"""Prompt construction — LangChain edition.
		2
		3	cortex-core built prompts with f-strings inside a `PromptBuilder` class. Here
		4	the same wording lives in LangChain prompt templates:
		5
		6	- `build_rag_prompt()` → a `ChatPromptTemplate` (system + human message) that
		7	the RAG graph pipes straight into the chat model: `prompt \| llm`.
		8	- `build_agentic_system_prompt()` → the agent's system prompt, rendered from a
		9	`PromptTemplate`.
		10
		11	A template is a Runnable, so it composes with `\|` and shows up as its own
		12	step in LangSmith/Langfuse traces. The wording is deliberately identical to
		13	cortex-core so the two implementations can be compared on equal footing.
		14	"""
		15
		16	from datetime import date
		17	from pathlib import Path
1	from typing import Dict, List, Optional	18	from typing import Dict, List, Optional
2	from datetime import date	19
3	import yaml	20	import yaml
4	from pathlib import Path	21	from langchain_core.documents import Document
		22	from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
		23
5		24
6	def _load_config() -> dict:	25	def _load_config() -> dict:
7	"""Load configuration from config.yml file."""	26	"""Load configuration from config.yml file."""
⋯ 2 unchanged lines hidden
8	# core/ -> business/ -> src/ -> config/config.yml	27	# core/ -> business/ -> src/ -> config/config.yml
9	config_path = Path(__file__).resolve().parents[2] / "config" / "config.yml"	28	config_path = Path(__file__).resolve().parents[2] / "config" / "config.yml"
10	with open(config_path, "r") as f:	29	with open(config_path, "r") as f:
11	return yaml.safe_load(f)	30	return yaml.safe_load(f)
12		31
		32
13	def _get_system_role() -> str:	33	def _get_system_role() -> str:
14	"""Get LLM system role from config.yml."""	34	"""Get LLM system role from config.yml."""
15	config = _load_config()	35	config = _load_config()
16	return config.get("llm_config", {}).get("llm_system_role", "You are a helpful assistant.")	36	return config.get("llm_config", {}).get("llm_system_role", "You are a helpful assistant.")
17		37
18	class PromptBuilder:	38
		39	# ---------------------------------------------------------------------------
		40	# RAG prompt
		41	# ---------------------------------------------------------------------------
		42	# Template variables: {system_prompt}, {context}, {question}. The grounding
		43	# rules are the hallucination guardrail of the RAG path — keep them verbatim.
		44	RAG_HUMAN_TEMPLATE = (
		45	"Use the following context to answer the question.\n"
		46	"{context}\n"
		47	"Question: {question}\n"
		48	"Rules:\n"
		49	"- Answer based only on the provided context.\n"
		50	'- If the answer is not in the context, respond with "I don\'t know."\n'
		51	"- Be concise and to the point.\n"
		52	"- Do not invent information.\n"
		53	)
		54
		55
		56	def build_rag_prompt(system_prompt: Optional[str] = None) -> ChatPromptTemplate:
		57	"""ChatPromptTemplate for the RAG generate step.
		58
		59	The system prompt is bound up front with .partial(), so the graph only has
		60	to supply {context} and {question} at run time.
		61
		62	Args:
		63	system_prompt: Custom system prompt. If None, uses llm_system_role from config.yml.
19	"""	64	"""
20	Responsible for constructing prompts for the AI model.	65	template = ChatPromptTemplate.from_messages([
21	Sent to the LLM to generate responses based on user input and context in RAG pipeline.	66	("system", "{system_prompt}"),
22	"""	67	("human", RAG_HUMAN_TEMPLATE),
		68	])
		69	return template.partial(system_prompt=system_prompt or _get_system_role())
23		70
24	def __init__(self, system_prompt: str = None):
25	"""
26	Initialize PromptBuilder.
27		71
28	Args:	72	def format_context(docs: List[Document]) -> str:
29	system_prompt: Custom system prompt. If None, uses llm_system_role from config.yml.	73	"""Render retrieved chunks as numbered context blocks: [Context 1]: ..."""
30	"""	74	return "\n\n".join(f"[Context {i + 1}]: {doc.page_content}" for i, doc in enumerate(docs))
31	# Use config value if system_prompt not provided
32	self.system_prompt = system_prompt or _get_system_role()
33		75
34	def build_prompt_text(self, question: str, context: List[str]) -> str:
35	"""
36	Constructs the full prompt (plain text) for completion-style models.
37	"""
38	context_block = "\n\n".join(f"[Context {i+1}]: {ctx}" for i, ctx in enumerate(context))
39	full_prompt = f"""
40	{self.system_prompt}
41	Use the following context to answer the question.
42	{context_block}
43	Question: {question}
44	Rules:
45	- Answer based only on the provided context.
46	- If the answer is not in the context, respond with "I don't know."
47	- Be concise and to the point.
48	- Do not invent information.
49	"""
50	return full_prompt.strip()
51		76
52	def build_messages(self, question: str, context: List[str]):	77	# ---------------------------------------------------------------------------
53	"""	78	# Agent system prompt
54	Constructs chat messages for OpenAI Chat Completions API.	79	# ---------------------------------------------------------------------------
55	"""	80	AGENTIC_SYSTEM_TEMPLATE = PromptTemplate.from_template("""You are a helpful, context-aware personal assistant.
56	context_block = "\n\n".join(f"[Context {i+1}]: {ctx}" for i, ctx in enumerate(context))
57	user_msg = (
58	"Use the following context to answer the question.\n"
59	f"{context_block}\n"
60	f"Question: {question}\n"
61	"Rules:\n"
62	'- Answer based only on the provided context.\n'
63	'- If the answer is not in the context, respond with "I don\'t know."\n'
64	"- Be concise and to the point.\n"
65	"- Do not invent information.\n"
66	)
67	return [
68	{"role": "system", "content": self.system_prompt},
69	{"role": "user", "content": user_msg},
70	]
71		81
72	# Backwards compatibility	82	Today's date is {today}. Your own training data is older than this, so treat
73	def build_prompt(self, question: str, context: List[str]) -> str:	83	anything time-sensitive as something you need to look up rather than recall.
74	return self.build_prompt_text(question, context)	84
		85	User profile:
		86	{user_block}
		87
		88	Conversation summary so far:
		89	{summary_block}
		90
		91	Relevant past conversations (semantic recall):
		92	{snippets}
		93
		94	Instructions:
		95	- Use the user profile and past context to personalise your responses.
		96	- Before answering, ALWAYS call search_vector_db to look up anything the user
		97	may have shared previously that is relevant to their current question — such
		98	as their job, hobbies, interests, preferences, or any personal details.
		99	Do not skip this step even if the current question seems unrelated to past
		100	conversations; the user's background often changes what a good answer looks like.
		101	- Call web_search for anything current: news, recent events, "latest" or
		102	"newest" anything, or facts that change over time. Search using the user's
		103	own terms — never add a year to the query yourself, as that biases the
		104	results toward the wrong period.
		105	- Results from web_search are more current than your training data. When they
		106	conflict, trust the search results and answer from them. If the results are
		107	dated later than you expect, they are still correct — do not dismiss them.
		108	- Be concise and direct. Do not repeat context back to the user verbatim.
		109	- If you are uncertain, say so rather than inventing information.""")
75		110
76		111
77	def build_agentic_system_prompt(	112	def build_agentic_system_prompt(
⋯ 2 unchanged lines hidden
78	user_info: Dict,	113	user_info: Dict,
79	vector_results: Optional[List[Dict]] = None,	114	vector_results: Optional[List[Dict]] = None,
80	chat_summary: Optional[str] = None,	115	chat_summary: Optional[str] = None,
81	) -> str:	116	) -> str:
82	"""	117	"""
83	Build the system prompt for the agentic chatbot (V3).	118	Build the system prompt for the agentic chatbot.
84		119
85	Injects user profile, an optional rolling summary, and semantically	120	Injects user profile, an optional rolling summary, and semantically
86	retrieved past-conversation snippets so the model has rich context	121	retrieved past-conversation snippets so the model has rich context
⋯ 22 unchanged lines hidden
87	without blowing up the context window with raw history.	122	without blowing up the context window with raw history.
88		123
89	Args:	124	Args:
90	user_info: Key/value pairs describing the user (name, preferences, etc.)	125	user_info: Key/value pairs describing the user (name, preferences, etc.)
91	vector_results: Output of LongTermMemory.recall() — list of	126	vector_results: Output of LongTermMemory.recall() — list of
92	{"text": str, "metadata": dict, "score": float}	127	{"text": str, "metadata": dict, "score": float}
93	chat_summary: Optional condensed summary of the conversation so far.	128	chat_summary: Optional condensed summary of the conversation so far.
94	"""	129	"""
95	user_block = "\n".join(f" {k}: {v}" for k, v in user_info.items()) if user_info else " (no user info)"	130	user_block = "\n".join(f" {k}: {v}" for k, v in user_info.items()) if user_info else " (no user info)"
96		131
97	summary_block = chat_summary.strip() if chat_summary else "(no summary yet)"	132	summary_block = chat_summary.strip() if chat_summary else "(no summary yet)"
98		133
99	if vector_results:	134	if vector_results:
100	snippets = "\n\n".join(	135	snippets = "\n\n".join(
101	f"[{i + 1}] {r['text']}" for i, r in enumerate(vector_results)	136	f"[{i + 1}] {r['text']}" for i, r in enumerate(vector_results)
102	)	137	)
103	else:	138	else:
104	snippets = "(no relevant past conversations found)"	139	snippets = "(no relevant past conversations found)"
105		140
106	# Without this the model falls back on its training cutoff as "now": it	141	# Without this the model falls back on its training cutoff as "now": it
107	# appends a stale year to web_search queries and then discounts the fresh	142	# appends a stale year to web_search queries and then discounts the fresh
108	# results it gets back as implausibly future-dated. Stating the real date	143	# results it gets back as implausibly future-dated. Stating the real date
109	# is what makes the web_search tool actually usable.	144	# is what makes the web_search tool actually usable.
110	today = date.today().strftime("%A, %d %B %Y")	145	today = date.today().strftime("%A, %d %B %Y")
111		146
112	return f"""You are a helpful, context-aware personal assistant.	147	return AGENTIC_SYSTEM_TEMPLATE.format(
113		148	today=today,
114	Today's date is {today}. Your own training data is older than this, so treat	149	user_block=user_block,
115	anything time-sensitive as something you need to look up rather than recall.	150	summary_block=summary_block,
116		151	snippets=snippets,
117	User profile:	152	)
118	{user_block}
119
120	Conversation summary so far:
121	{summary_block}
122
123	Relevant past conversations (semantic recall):
124	{snippets}
125
126	Instructions:
127	- Use the user profile and past context to personalise your responses.
128	- Before answering, ALWAYS call search_vector_db to look up anything the user
129	may have shared previously that is relevant to their current question — such
130	as their job, hobbies, interests, preferences, or any personal details.
131	Do not skip this step even if the current question seems unrelated to past
132	conversations; the user's background often changes what a good answer looks like.
133	- Call web_search for anything current: news, recent events, "latest" or
134	"newest" anything, or facts that change over time. Search using the user's
135	own terms — never add a year to the query yourself, as that biases the
136	results toward the wrong period.
137	- Results from web_search are more current than your training data. When they
138	conflict, trust the search results and answer from them. If the results are
139	dated later than you expect, they are still correct — do not dismiss them.
140	- Be concise and direct. Do not repeat context back to the user verbatim.
141	- If you are uncertain, say so rather than inventing information."""

07 retrieval.py

src/business/rag/retrieval.py · 78 → 182 lines · 159 rows changed

The RAG pipeline: retrieve, re-rank, then answer. cortex-core ran it as three calls in a row; cortex-langgraph makes each step a graph node.

cortex-core · hand-written

RAGPipeline.answer() embeds the question, queries the store, re-ranks, then calls the model.
_retrieve() parses Chroma's raw nested-list response into RetrievedChunk objects.
What happens when the re-ranker rejects everything is buried inside select_context() in orchestrator.py.

cortex-langgraph · LangChain / LangGraph

A StateGraph: retrieve → rerank → generate, with fallback and refuse branches.
The fallback policy (fail_closed / fail_open / hybrid) is a visible conditional edge, _route_after_rerank().
generate is an LCEL chain: ChatPromptTemplate | chat model | StrOutputParser.
State records confidence (high / low / none) to show which branch ran.

Why it matters: The safety decision ("refuse or fall back when nothing relevant was found") is now part of the graph you can see, instead of hidden in a helper function.

StateGraphconditional edgesStrOutputParsersimilarity_search_with_score

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/retrieval.py		cortex-langgraph src/business/rag/retrieval.py
1	"""Retrieve → rerank → prompt → generate."""	1	"""RAG pipeline — a LangGraph StateGraph.
		2
		3	cortex-core's `RAGPipeline.answer()` was three function calls in a row, with
		4	the "what if the re-ranker rejects everything?" policy buried inside
		5	`select_context()`. Here each step is a node and that policy is a visible
		6	branch:
		7
		8	START → retrieve → rerank ──(chunks passed the gate)──────────────▶ generate → END
		9	│ ▲
		10	├──(nothing passed, policy fail_open/hybrid)──▶ fallback
		11	│
		12	└──(nothing passed, policy fail_closed)───────▶ refuse → END
		13
		14	- retrieve VectorStore.similarity_search_with_score — dense top-k (default 30)
		15	- rerank cross-encoder score → relevance gate → sort → top-n (default 8)
		16	- fallback gate rejected everything: use the raw vector top-n, confidence "low"
		17	- refuse fail_closed policy: don't call the LLM at all
		18	- generate ChatPromptTemplate \| chat model \| StrOutputParser (an LCEL chain)
		19
		20	Every node reads the shared state and returns only the keys it changes; LangGraph
		21	merges them. `confidence` records which branch ran — "high" (re-ranked),
		22	"low" (fallback) or "none" (refused) — and query_rag() turns it into metrics.
		23	"""
2		24
3	from __future__ import annotations	25	from __future__ import annotations
4		26
5	from typing import List, Tuple	27	from typing import List, Literal, Optional, Tuple, TypedDict
6		28
7	from dotenv import load_dotenv	29	from dotenv import load_dotenv
		30	from langchain_core.cross_encoders import BaseCrossEncoder
		31	from langchain_core.documents import Document
		32	from langchain_core.language_models import BaseChatModel
		33	from langchain_core.output_parsers import StrOutputParser
		34	from langchain_core.vectorstores import VectorStore
		35	from langgraph.graph import END, START, StateGraph
8		36
9	from ..core.embedding import create_embedder	37	from ..core.model import create_llm
		38	from ..core.prompt_builder import build_rag_prompt, format_context
		39	from .re_ranker.config import ReRankerConfig
		40	from .re_ranker.re_ranker import rerank_documents
10	from .vector_store import create_vector_store	41	from .vector_store import create_vector_store
11	from .re_ranker.interface import RetrievedChunk, ReRankedChunk	42
12	from .re_ranker.re_ranker import ReRanker	43	Policy = Literal["fail_closed", "fail_open", "hybrid"]
13	from .re_ranker.config import ReRankerConfig	44
14	from .re_ranker.cross_encoder import CrossEncoderReRanker	45	# What the refuse branch answers — the same wording the grounding rules tell
15	from .re_ranker.orchestrator import select_context	46	# the model to use when the context doesn't contain the answer.
16	from ..core.prompt_builder import PromptBuilder	47	REFUSAL = "I don't know."
17	from ..core.model import create_llm	48
		49
		50	class RAGState(TypedDict, total=False):
		51	"""The data that flows through the graph. Nodes return partial updates."""
		52	question: str
		53	candidates: List[Document] # vector-store results, best first
		54	context: List[Document] # what the LLM will see
		55	confidence: str # "high" \| "low" \| "none"
		56	answer: str
18		57
19		58
20	class RAGPipeline:	59	class RAGPipeline:
		60	"""
		61	Retrieve → rerank → (gate) → generate, compiled as a LangGraph graph.
		62
		63	Collaborators are built from the provider factories by default, and can be
		64	injected (tests pass fakes; evals pass a shared cross-encoder).
		65	"""
		66
21	def __init__(	67	def __init__(
22	self,	68	self,
23	persist_dir: str,	69	persist_dir: str,
⋯ 1 unchanged line hidden
24	collection_name: str = "pdf_chunks",	70	collection_name: str = "pdf_chunks",
25	reranker_config: ReRankerConfig \| None = None,	71	reranker_config: ReRankerConfig \| None = None,
26	system_prompt: str \| None = None,	72	system_prompt: str \| None = None,
27	model_name: str = "gpt-4o-mini",	73	model_name: str = "gpt-4o-mini",
		74	policy: Policy = "hybrid",
		75	vector_store: Optional[VectorStore] = None,
		76	scorer: Optional[BaseCrossEncoder] = None,
		77	llm: Optional[BaseChatModel] = None,
28	):	78	):
29	load_dotenv()	79	load_dotenv()
		80
		81	self.reranker_config = reranker_config or ReRankerConfig()
		82	self.policy = policy
30		83
31	# Each factory reads its own *_PROVIDER env var (defaulting to	84	# Each factory reads its own *_PROVIDER env var (defaulting to
32	# today's OpenAI-based behavior) and only requires the credentials	85	# today's OpenAI-based behavior) and only requires the credentials
33	# its selected provider actually needs.	86	# its selected provider actually needs.
34	self.embedder = create_embedder()	87	self.vector_store = vector_store or create_vector_store(
35	self.vector_store = create_vector_store(persist_dir=persist_dir, collection_name=collection_name)	88	persist_dir=persist_dir, collection_name=collection_name
		89	)
		90	if scorer is None:
		91	# Imported lazily: loading the cross-encoder pulls in torch.
		92	from .re_ranker.cross_encoder import CrossEncoderScorer
36		93
37	scorer = CrossEncoderReRanker(reranker_config or ReRankerConfig())	94	scorer = CrossEncoderScorer(self.reranker_config)
38	self.reranker = ReRanker(scorer=scorer, config=reranker_config or ReRankerConfig())	95	self.scorer = scorer
		96	# Same sampling settings cortex-core's OpenAIModel defaulted to.
		97	self.llm = llm or create_llm(model_name=model_name, temperature=0.7, max_tokens=512)
39		98
40	self.reranker_config = reranker_config or ReRankerConfig()	99	# prompt \| llm \| parser — an LCEL chain, used by the generate node.
41	self.prompt_builder = PromptBuilder(system_prompt)	100	self.generate_chain = build_rag_prompt(system_prompt) \| self.llm \| StrOutputParser()
42	self.llm = create_llm(model_name=model_name, system_prompt=system_prompt)	101	self.graph = self._build_graph()
43		102
44	def _retrieve(self, query: str, top_k: int = 30) -> List[RetrievedChunk]:	103	# ------------------------------------------------------------------ nodes
45	q_emb = self.embedder.embed_query(query)	104	def _retrieve(self, state: RAGState) -> RAGState:
46	result = self.vector_store.query(q_emb, top_k)	105	results = self.vector_store.similarity_search_with_score(
47	retrieved: List[RetrievedChunk] = []	106	state["question"], k=self.reranker_config.top_k_input
48	ids = result.get("ids", [[]])[0]	107	)
49	docs = result.get("documents", [[]])[0]	108	candidates = [
50	metas = result.get("metadatas", [[]])[0]	109	Document(
51	dists = result.get("distances", [[]])[0]	110	id=doc.id,
52	for chunk_id, doc, meta, dist in zip(ids, docs, metas, dists):	111	page_content=doc.page_content,
53	retrieved.append(	112	metadata={**doc.metadata, "vector_score": float(score)},
54	RetrievedChunk(
55	chunk_id=chunk_id,
56	text=doc,
57	metadata=meta or {},
58	vector_score=float(dist) if dist is not None else 0.0,
59	)
60	)	113	)
61	return retrieved	114	for doc, score in results
62	# The answer function is the orchestration of the RAG pipeline not just the retrieval step.	115	]
63	def answer(self, question: str) -> Tuple[str, List[ReRankedChunk]]:	116	return {"candidates": candidates}
64	# Step 1: Retrieve relevant chunks from vector store
65	retrieved_chunks = self._retrieve(question, top_k=self.reranker_config.top_k_input)
66		117
67	# Step 2: Re-rank and select context (orchestrator handles reranking internally)	118	def _rerank(self, state: RAGState) -> RAGState:
68	selected_context, confidence = select_context(	119	reranked = rerank_documents(
69	query=question,	120	query=state["question"],
70	retrieved_chunks=retrieved_chunks,	121	documents=state["candidates"],
71	reranker=self.reranker,	122	scorer=self.scorer,
72	policy="hybrid",	123	config=self.reranker_config,
73	)	124	)
		125	# An empty context tells the router below that the gate rejected everything.
		126	return {"context": reranked, "confidence": "high" if reranked else "none"}
74		127
75	# Step 3: Generate answer using LLM	128	def _route_after_rerank(self, state: RAGState) -> str:
76	answer_text = self.llm.generate(question, [chunk.text for chunk in selected_context])	129	"""The policy cortex-core kept in orchestrator.select_context()."""
		130	if state["context"]:
		131	return "generate"
		132	if self.policy == "fail_closed":
		133	return "refuse"
		134	return "fallback"
77		135
78	return answer_text, selected_context	136	def _fallback(self, state: RAGState) -> RAGState:
		137	# Fail open: raw vector order, flagged low-confidence so it shows up in metrics.
		138	return {
		139	"context": state["candidates"][: self.reranker_config.top_n_output],
		140	"confidence": "low",
		141	}
		142
		143	def _refuse(self, state: RAGState) -> RAGState:
		144	return {"context": [], "confidence": "none", "answer": REFUSAL}
		145
		146	def _generate(self, state: RAGState) -> RAGState:
		147	answer = self.generate_chain.invoke({
		148	"question": state["question"],
		149	"context": format_context(state["context"]),
		150	})
		151	return {"answer": answer.strip()}
		152
		153	# ------------------------------------------------------------------ graph
		154	def _build_graph(self):
		155	builder = StateGraph(RAGState)
		156	builder.add_node("retrieve", self._retrieve)
		157	builder.add_node("rerank", self._rerank)
		158	builder.add_node("fallback", self._fallback)
		159	builder.add_node("refuse", self._refuse)
		160	builder.add_node("generate", self._generate)
		161
		162	builder.add_edge(START, "retrieve")
		163	builder.add_edge("retrieve", "rerank")
		164	builder.add_conditional_edges(
		165	"rerank",
		166	self._route_after_rerank,
		167	{"generate": "generate", "fallback": "fallback", "refuse": "refuse"},
		168	)
		169	builder.add_edge("fallback", "generate")
		170	builder.add_edge("generate", END)
		171	builder.add_edge("refuse", END)
		172	return builder.compile()
		173
		174	# -------------------------------------------------------------------- API
		175	def run(self, question: str) -> RAGState:
		176	"""Run the graph and return its final state."""
		177	return self.graph.invoke({"question": question})
		178
		179	def answer(self, question: str) -> Tuple[str, List[Document], str]:
		180	"""Answer a question. Returns (answer, context chunks, confidence)."""
		181	state = self.run(question)
		182	return state["answer"], state.get("context", []), state.get("confidence", "none")

08 vector_store.py

src/business/rag/vector_store.py · 261 → 147 lines · 316 rows changed

Where PDF chunks are stored and searched. The biggest reduction in code: LangChain's VectorStore interface replaces two hand-written backends.

cortex-core · hand-written

A hand-written VectorStoreBase (reset / upsert / query) with two backends: ChromaVectorStore and AzureSearchVectorStore.
The Azure backend translates Azure Search results into Chroma's nested-list shape, because retrieval.py parsed exactly that shape.
Callers embed text first and pass in raw vectors.

cortex-langgraph · LangChain / LangGraph

create_vector_store() returns a LangChain VectorStore: langchain_chroma.Chroma or AzureSearch.
Both return (Document, score) pairs from similarity_search_with_score(), so no translation layer is needed.
The store owns its embedding model and embeds text itself.
Uses a different Azure index name, because LangChain's index schema isn't compatible with cortex-core's.

Why it matters: Swapping backends no longer means writing a new adapter class and a response translator.

VectorStoreChromaAzureSearchDocument

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/vector_store.py		cortex-langgraph src/business/rag/vector_store.py
1	"""Chroma vector store wrapper.	1	"""RAG vector store factory — LangChain edition.
2		2
3	For production RAG/chatbot systems you may prefer managed vector DBs	3	cortex-core defined its own `VectorStoreBase` (reset/upsert/query) with two
4	like Weaviate or Pinecone. Chroma is used here for simplicity; swap it	4	hand-written backends, and the Azure one had to translate Azure Search's
5	out for scalability/performance in production.	5	response into Chroma's nested-list dict shape because retrieval.py parsed that
		6	exact structure.
		7
		8	LangChain's `VectorStore` is a common interface, so none of that is needed:
		9	both backends return `(Document, score)` pairs from
		10	`similarity_search_with_score()`, and they embed text themselves given an
		11	`Embeddings` object.
		12
		13	For production RAG systems you may prefer a managed vector DB. Chroma is the
		14	default for simplicity; Azure AI Search is the cloud alternative.
6	"""	15	"""
7		16
8	from __future__ import annotations	17	from __future__ import annotations
9		18
10	import os	19	import os
11	from abc import ABC, abstractmethod	20	from typing import Optional
12	from typing import List, Dict, Any, Optional
13	import chromadb
14	from chromadb.utils import embedding_functions
15		21
		22	from chromadb.config import Settings
		23	from langchain_core.embeddings import Embeddings
		24	from langchain_core.vectorstores import VectorStore
16		25
17	class VectorStoreBase(ABC):	26	DEFAULT_COLLECTION = "pdf_chunks"
18	"""Interface every RAG vector store backend must implement."""
19		27
20	@abstractmethod	28	# Different default from cortex-core's "rag-chunks": LangChain's AzureSearch
21	def reset(self) -> None: ...	29	# uses its own index schema (id, content, content_vector, metadata), which is
22		30	# not compatible with the index cortex-core creates. Distinct names mean the
23	@abstractmethod	31	# two projects can share one Azure Search service without clobbering each other.
24	def upsert(	32	DEFAULT_AZURE_INDEX = "rag-chunks-langchain"
25	self,
26	ids: List[str],
27	embeddings: List[List[float]],
28	metadatas: List[Dict[str, Any]],
29	documents: List[str],
30	) -> None: ...
31
32	@abstractmethod
33	def query(self, query_embedding: List[float], top_k: int = 15): ...
34
35
36	class ChromaVectorStore(VectorStoreBase):
37	def __init__(self, persist_dir: str, collection_name: str = "pdf_chunks", dim: int \| None = None):
38	self.client = chromadb.PersistentClient(path=persist_dir)
39	self.collection = self.client.get_or_create_collection(
40	name=collection_name,
41	embedding_function=None, # we supply embeddings manually
42	)
43	self.dim = dim
44
45	def reset(self):
46	self.client.delete_collection(self.collection.name)
47	self.collection = self.client.get_or_create_collection(self.collection.name)
48
49	def upsert(
50	self,
51	ids: List[str],
52	embeddings: List[List[float]],
53	metadatas: List[Dict[str, Any]],
54	documents: List[str],
55	):
56	# check the lengths matches for ids, embeddings, metadatas, documents
57	if len(ids) != len(embeddings):
58	raise ValueError("ids and embeddings length mismatch")
59	self.collection.upsert(
60	ids=ids,
61	embeddings=embeddings,
62	metadatas=metadatas,
63	documents=documents,
64	)
65	# Run the similarity search and return top_k results
66	def query(self, query_embedding: List[float], top_k: int = 15):
67	return self.collection.query(
68	query_embeddings=[query_embedding],
69	n_results=top_k,
70	include=["documents", "metadatas", "distances"],
71	)
72
73
74	# Backward-compat alias — existing code imports `VectorStore` directly.
75	VectorStore = ChromaVectorStore
76
77
78	class AzureSearchVectorStore(VectorStoreBase):
79	"""Azure AI Search-backed vector store for RAG chunks — same interface as
80	ChromaVectorStore, backed by a cloud vector index instead of local Chroma.
81
82	Creates the index on first use if it doesn't already exist yet. reset()
83	drops and recreates the index, matching ChromaVectorStore.reset()'s
84	"wipe to an empty collection" semantics exactly.
85
86	The `azure-search-documents` SDK is imported lazily (inside methods, not
87	at module level) so on-prem/Chroma-only deployments never need it installed.
88	"""
89
90	INDEX_FIELDS_METADATA_KEYS = ("source_id", "section", "chunk_start", "chunk_end", "chunk_strategy")
91
92	def __init__(
93	self,
94	endpoint: str,
95	api_key: str,
96	index_name: str = "rag-chunks",
97	dim: int = 1536,
98	):
99	from azure.core.credentials import AzureKeyCredential
100	from azure.core.exceptions import ResourceNotFoundError
101	from azure.search.documents import SearchClient
102	from azure.search.documents.indexes import SearchIndexClient
103	from azure.search.documents.indexes.models import (
104	HnswAlgorithmConfiguration,
105	SearchField,
106	SearchFieldDataType,
107	SearchIndex,
108	SimpleField,
109	VectorSearch,
110	VectorSearchProfile,
111	)
112
113	self._ResourceNotFoundError = ResourceNotFoundError
114	self._SearchIndex = SearchIndex
115	self.endpoint = endpoint
116	self.index_name = index_name
117	self.dim = dim
118
119	credential = AzureKeyCredential(api_key)
120	self._index_client = SearchIndexClient(endpoint=endpoint, credential=credential)
121	self._search_client = SearchClient(endpoint=endpoint, index_name=index_name, credential=credential)
122
123	algorithm_name = "rag-hnsw"
124	profile_name = "rag-vector-profile"
125	self._vector_search = VectorSearch(
126	algorithms=[HnswAlgorithmConfiguration(name=algorithm_name)],
127	profiles=[VectorSearchProfile(name=profile_name, algorithm_configuration_name=algorithm_name)],
128	)
129	self._fields = [
130	SimpleField(name="id", type=SearchFieldDataType.STRING, key=True),
131	SearchField(name="content", type=SearchFieldDataType.STRING, searchable=True),
132	SearchField(
133	name="embedding",
134	type="Collection(Edm.Single)",
135	searchable=True,
136	vector_search_dimensions=dim,
137	vector_search_profile_name=profile_name,
138	),
139	SimpleField(name="source_id", type=SearchFieldDataType.STRING, filterable=True),
140	SimpleField(name="section", type=SearchFieldDataType.STRING, filterable=True),
141	SimpleField(name="chunk_start", type=SearchFieldDataType.INT32, filterable=True),
142	SimpleField(name="chunk_end", type=SearchFieldDataType.INT32, filterable=True),
143	SimpleField(name="chunk_strategy", type=SearchFieldDataType.STRING, filterable=True),
144	]
145
146	self._ensure_index_exists()
147
148	def _ensure_index_exists(self) -> None:
149	try:
150	self._index_client.get_index(self.index_name)
151	except self._ResourceNotFoundError:
152	self._index_client.create_index(
153	self._SearchIndex(name=self.index_name, fields=self._fields, vector_search=self._vector_search)
154	)
155
156	def reset(self) -> None:
157	try:
158	self._index_client.delete_index(self.index_name)
159	except self._ResourceNotFoundError:
160	pass
161	self._index_client.create_index(
162	self._SearchIndex(name=self.index_name, fields=self._fields, vector_search=self._vector_search)
163	)
164
165	def upsert(
166	self,
167	ids: List[str],
168	embeddings: List[List[float]],
169	metadatas: List[Dict[str, Any]],
170	documents: List[str],
171	) -> None:
172	if len(ids) != len(embeddings):
173	raise ValueError("ids and embeddings length mismatch")
174
175	docs = []
176	for doc_id, embedding, metadata, text in zip(ids, embeddings, metadatas, documents):
177	metadata = metadata or {}
178	docs.append({
179	"id": doc_id,
180	"content": text,
181	"embedding": embedding,
182	"source_id": str(metadata.get("source_id", "")),
183	"section": str(metadata.get("section", "")),
184	"chunk_start": int(metadata.get("chunk_start", 0)),
185	"chunk_end": int(metadata.get("chunk_end", 0)),
186	"chunk_strategy": str(metadata.get("chunk_strategy", "")),
187	})
188	self._search_client.merge_or_upload_documents(documents=docs)
189
190	def query(self, query_embedding: List[float], top_k: int = 15):
191	from azure.search.documents.models import VectorizedQuery
192
193	vector_query = VectorizedQuery(vector=query_embedding, k_nearest_neighbors=top_k, fields="embedding")
194	results = self._search_client.search(search_text=None, vector_queries=[vector_query], top=top_k)
195
196	ids: List[str] = []
197	docs: List[str] = []
198	metas: List[Dict[str, Any]] = []
199	scores: List[float] = []
200	for r in results:
201	ids.append(r["id"])
202	docs.append(r.get("content", ""))
203	metas.append({key: r.get(key) for key in self.INDEX_FIELDS_METADATA_KEYS})
204	# Azure Search's vector-query score is a similarity score (higher =
205	# better); Chroma's "distances" are a distance (lower = better).
206	# Negate it so "lower is better" holds for both backends — nothing
207	# downstream treats this as a calibrated distance, it's only ever
208	# used for display and (already-sorted) ordering.
209	scores.append(-float(r.get("@search.score", 0.0)))
210
211	# Match ChromaVectorStore.query()'s response shape exactly (outer list
212	# = batch of queries; we only ever send one) since retrieval.py's
213	# _retrieve() depends on this exact structure regardless of backend.
214	return {
215	"ids": [ids],
216	"documents": [docs],
217	"metadatas": [metas],
218	"distances": [scores],
219	}
220		33
221		34
222	# ── Provider selection ────────────────────────────────────────────────────	35	# ── Provider selection ────────────────────────────────────────────────────
⋯ 2 unchanged lines hidden
223	#	36	#
224	# VECTOR_STORE_PROVIDER (env var) picks the RAG vector store backend.	37	# VECTOR_STORE_PROVIDER (env var) picks the RAG vector store backend.
225	# Defaults to today's behavior — on-prem/local deployments don't need to	38	# Defaults to today's behavior — on-prem/local deployments don't need to
226	# set anything.	39	# set anything.
227	#	40	#
228	# chroma (default) — local Chroma persistence.	41	# chroma (default) — langchain_chroma.Chroma, local persistence.
229	# azure_search — Azure AI Search (vector search). Requires	42	# azure_search — langchain_azure_ai AzureSearch. Requires
230	# AZURE_SEARCH_ENDPOINT, AZURE_SEARCH_API_KEY, and	43	# AZURE_SEARCH_ENDPOINT, AZURE_SEARCH_API_KEY, and
231	# optionally AZURE_SEARCH_INDEX_NAME (see .env).	44	# optionally AZURE_SEARCH_INDEX_NAME / AZURE_SEARCH_EMBEDDING_DIM.
		45
		46	def _provider(provider: Optional[str]) -> str:
		47	return (provider or os.getenv("VECTOR_STORE_PROVIDER", "chroma")).strip().lower()
		48
232		49
233	def create_vector_store(	50	def create_vector_store(
234	persist_dir: str,	51	persist_dir: str,
235	collection_name: str = "pdf_chunks",	52	collection_name: str = DEFAULT_COLLECTION,
236	provider: Optional[str] = None,	53	provider: Optional[str] = None,
237	dim: Optional[int] = None,	54	embedding: Optional[Embeddings] = None,
238	) -> VectorStoreBase:	55	) -> VectorStore:
239	"""Factory for the RAG vector store, selected by VECTOR_STORE_PROVIDER or `provider`."""	56	"""Factory for the RAG vector store, selected by VECTOR_STORE_PROVIDER or `provider`.
240	provider = (provider or os.getenv("VECTOR_STORE_PROVIDER", "chroma")).strip().lower()	57
		58	embedding: the model used to embed chunks and queries. Defaults to
		59	create_embedder() — it must be the same model at index and query time.
		60	"""
		61	provider = _provider(provider)
		62	if provider not in ("chroma", "azure_search"):
		63	raise ValueError(
		64	f"Unknown VECTOR_STORE_PROVIDER={provider!r}. Supported: chroma, azure_search."
		65	)
		66
		67	if embedding is None:
		68	from src.business.core.embedding import create_embedder
		69
		70	embedding = create_embedder()
241		71
242	if provider == "chroma":	72	if provider == "chroma":
243	return ChromaVectorStore(persist_dir=persist_dir, collection_name=collection_name)	73	return _chroma(persist_dir, collection_name, embedding)
		74
		75	endpoint, api_key, index_name = _azure_settings()
		76	# Imported lazily so Chroma-only deployments never load the Azure SDK.
		77	from langchain_azure_ai.vectorstores import AzureSearch
		78
		79	return AzureSearch(
		80	azure_search_endpoint=endpoint,
		81	azure_search_key=api_key,
		82	index_name=index_name,
		83	embedding_function=embedding,
		84	# Passing the dimension up front stops AzureSearch from embedding a
		85	# probe string (a paid API call) just to discover it. The index is
		86	# created on first use if it doesn't exist.
		87	vector_search_dimensions=int(os.getenv("AZURE_SEARCH_EMBEDDING_DIM", "1536")),
		88	)
		89
		90
		91	def reset_vector_store(
		92	persist_dir: str,
		93	collection_name: str = DEFAULT_COLLECTION,
		94	provider: Optional[str] = None,
		95	) -> None:
		96	"""Wipe the RAG index to empty.
		97
		98	Upserts never delete, so stale chunks from a previous upload would
		99	otherwise keep matching queries. Needs no embedding model.
		100	"""
		101	provider = _provider(provider)
		102
		103	if provider == "chroma":
		104	_chroma(persist_dir, collection_name, embedding=None).reset_collection()
		105	return
244		106
245	if provider == "azure_search":	107	if provider == "azure_search":
246	endpoint = os.getenv("AZURE_SEARCH_ENDPOINT")	108	endpoint, api_key, index_name = _azure_settings()
247	api_key = os.getenv("AZURE_SEARCH_API_KEY")	109	from azure.core.credentials import AzureKeyCredential
248	index_name = os.getenv("AZURE_SEARCH_INDEX_NAME", "rag-chunks")	110	from azure.core.exceptions import ResourceNotFoundError
249	if not endpoint or not api_key:	111	from azure.search.documents.indexes import SearchIndexClient
250	raise RuntimeError(	112
251	"VECTOR_STORE_PROVIDER=azure_search requires AZURE_SEARCH_ENDPOINT "	113	client = SearchIndexClient(endpoint=endpoint, credential=AzureKeyCredential(api_key))
252	"and AZURE_SEARCH_API_KEY."	114	try:
253	)	115	# AzureSearch recreates the index on its next construction.
254	resolved_dim = dim or int(os.getenv("AZURE_SEARCH_EMBEDDING_DIM", "1536"))	116	client.delete_index(index_name)
255	return AzureSearchVectorStore(	117	except ResourceNotFoundError:
256	endpoint=endpoint, api_key=api_key, index_name=index_name, dim=resolved_dim	118	pass
257	)	119	return
258		120
259	raise ValueError(	121	raise ValueError(
260	f"Unknown VECTOR_STORE_PROVIDER={provider!r}. Supported: chroma, azure_search."	122	f"Unknown VECTOR_STORE_PROVIDER={provider!r}. Supported: chroma, azure_search."
261	)	123	)
		124
		125
		126	def _chroma(persist_dir: str, collection_name: str, embedding: Optional[Embeddings]) -> VectorStore:
		127	from langchain_chroma import Chroma
		128
		129	return Chroma(
		130	collection_name=collection_name,
		131	embedding_function=embedding,
		132	persist_directory=persist_dir,
		133	client_settings=Settings(anonymized_telemetry=False, is_persistent=True,
		134	persist_directory=persist_dir),
		135	)
		136
		137
		138	def _azure_settings() -> tuple[str, str, str]:
		139	endpoint = os.getenv("AZURE_SEARCH_ENDPOINT")
		140	api_key = os.getenv("AZURE_SEARCH_API_KEY")
		141	index_name = os.getenv("AZURE_SEARCH_INDEX_NAME", DEFAULT_AZURE_INDEX)
		142	if not endpoint or not api_key:
		143	raise RuntimeError(
		144	"VECTOR_STORE_PROVIDER=azure_search requires AZURE_SEARCH_ENDPOINT "
		145	"and AZURE_SEARCH_API_KEY."
		146	)
		147	return endpoint, api_key, index_name

09 index_builder.py

src/business/rag/index_builder.py · 97 → 70 lines · 59 rows changed

Building the search index from PDFs: load, split, store. Each step is now a LangChain component.

cortex-core · hand-written

ingest_directory() → a custom Chunker per document → embed each batch → store.upsert() with raw vectors.

cortex-langgraph · LangChain / LangGraph

DoclingLoader → RecursiveCharacterTextSplitter → VectorStore.add_documents().
add_documents() embeds and upserts in one call.

Why it matters: The pipeline shrinks to three calls, and every stage speaks the same Document type.

DoclingLoaderRecursiveCharacterTextSplitteradd_documents

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/index_builder.py		cortex-langgraph src/business/rag/index_builder.py
1	"""Index builder:	1	"""Index builder — LangChain edition.
2	Ingests PDFs → extracts text and tables	2
3	Chunks → splits into smaller pieces	3	Load → split → store, each step a LangChain component:
4	Embeds → converts text to vectors	4
5	Stores → saves to vector database (line 78)	5	DoclingLoader PDFs → Documents (pdfingest/pdf_digest.py)
		6	RecursiveCharacterTextSplitter Documents → chunks (pdfingest/chunk.py)
		7	VectorStore.add_documents chunks → embedded + stored (vector_store.py)
		8
		9	cortex-core embedded each batch itself and handed raw vectors to the store.
		10	A LangChain vector store owns its embedding model, so add_documents() embeds
		11	and upserts in one call.
6	"""	12	"""
7		13
8	from __future__ import annotations	14	from __future__ import annotations
9		15
10	from pathlib import Path	16	from pathlib import Path
11	from typing import List, Tuple	17	from typing import Tuple
12		18
13	from dotenv import load_dotenv	19	from dotenv import load_dotenv
14		20
15	from .pdfingest.pdf_digest import ingest_directory, IngestedDocument	21	from .pdfingest.pdf_digest import load_pdf_documents
16	from .pdfingest.chunk import Chunker, Chunk	22	from .pdfingest.chunk import split_documents
17	from ..core.embedding import create_embedder
18	from .vector_store import create_vector_store	23	from .vector_store import create_vector_store
19
20
21	def _chunk_document(doc: IngestedDocument, chunker: Chunker, table_start_idx: int \| None) -> List[Chunk]:
22	source_metadata = {"source_id": doc.source_id}
23	chunks = chunker.split(doc.combined_text, source_metadata=source_metadata)
24	enriched = []
25	for ch in chunks:
26	md = dict(ch.metadata)
27	section = "table" if table_start_idx is not None and md.get("chunk_start", 0) >= table_start_idx else "text"
28	md.update({"source_id": doc.source_id, "section": section})
29	enriched.append(Chunk(chunk_id=ch.chunk_id, text=ch.text, metadata=md))
30	return enriched
31		24
32		25
33	def build_index(	26	def build_index(
⋯ 12 unchanged lines hidden
34	data_dir: Path,	27	data_dir: Path,
35	persist_dir: Path,	28	persist_dir: Path,
36	max_context_chars: int = 12_000,	29	max_context_chars: int = 12_000,
37	include_table_images: bool = True,	30	include_table_images: bool = True,
38	pdf_strategy: str = "hi_res",	31	pdf_strategy: str = "hi_res",
39	chunk_size: int = 800,	32	chunk_size: int = 800,
40	overlap: int = 100,	33	overlap: int = 100,
41	top_k_store: int \| None = None,	34	top_k_store: int \| None = None,
42	) -> Tuple[int, int]:	35	) -> Tuple[int, int]:
43	"""	36	"""
44	Build or rebuild the vector index from PDFs.	37	Build or rebuild the vector index from PDFs.
45	Returns (docs_indexed, chunks_indexed).	38	Returns (docs_indexed, chunks_indexed).
46	"""	39	"""
47	load_dotenv()	40	load_dotenv()
48		41
49	docs = ingest_directory(	42	docs = load_pdf_documents(
50	data_dir=data_dir,	43	data_dir=data_dir,
51	max_context_chars=max_context_chars,	44	max_context_chars=max_context_chars,
52	include_table_images=include_table_images,	45	include_table_images=include_table_images,
⋯ 1 unchanged line hidden
53	pdf_strategy=pdf_strategy,	46	pdf_strategy=pdf_strategy,
54	)	47	)
55	print(f"📄 Chunking {len(docs)} document(s)...", flush=True)	48	print(f"📄 Chunking {len(docs)} document(s)...", flush=True)
56		49
57	embedder = create_embedder()	50	chunks = split_documents(docs, chunk_size=chunk_size, overlap=overlap)
58	chunker = Chunker(chunk_size=chunk_size, overlap=overlap)
59
60	all_chunks: List[Chunk] = []
61	for doc in docs:
62	print(f" Chunking doc: {doc.source_id} ({len(doc.combined_text)} chars)...", flush=True)
63	# use marker to identify start of table section
64	table_marker = None
65	if doc.table_text:
66	# compute start idx of table section for provenance tagging
67	marker = "\n\n" + "=" * 80 + "\nTABLES FROM DOCUMENT IMAGES:\n" + "=" * 80 + "\n"
68	table_marker = doc.combined_text.find(marker)
69	if table_marker == -1:
70	table_marker = None
71	doc_chunks = _chunk_document(doc, chunker, table_start_idx=table_marker)
72	all_chunks.extend(doc_chunks)
73
74	if top_k_store:	51	if top_k_store:
75	all_chunks = all_chunks[:top_k_store]	52	chunks = chunks[:top_k_store]
76		53	print(f"✅ Chunking complete: {len(chunks)} chunk(s)", flush=True)
77	print(f"✅ Chunking complete: {len(all_chunks)} chunk(s)", flush=True)
78		54
79	# Process in batches to avoid OOM when embedding/upserting many chunks	55	# Process in batches to avoid OOM when embedding/upserting many chunks
80	batch_size = 50	56	batch_size = 50
81	total_batches = (len(all_chunks) + batch_size - 1) // batch_size	57	total_batches = (len(chunks) + batch_size - 1) // batch_size
82	vstore = None	58	vstore = None
83	for i in range(0, len(all_chunks), batch_size):	59	for i in range(0, len(chunks), batch_size):
84	batch_num = i // batch_size + 1	60	batch_num = i // batch_size + 1
85	print(f"📤 Embedding batch {batch_num}/{total_batches}...", flush=True)
86	batch = all_chunks[i : i + batch_size]
87	texts = [c.text for c in batch]
88	embeddings = embedder.embed_documents(texts)
89	if vstore is None:	61	if vstore is None:
90	print("🔗 Connecting to vector store...", flush=True)	62	print("🔗 Connecting to vector store...", flush=True)
91	vstore = create_vector_store(persist_dir=str(persist_dir))	63	vstore = create_vector_store(persist_dir=str(persist_dir))
92	ids = [c.chunk_id for c in batch]	64	print(f"📤 Embedding + storing batch {batch_num}/{total_batches}...", flush=True)
93	metadatas = [c.metadata for c in batch]	65	batch = chunks[i : i + batch_size]
94	vstore.upsert(ids=ids, embeddings=embeddings, metadatas=metadatas, documents=texts)	66	# Explicit ids = the stable chunk ids, so re-indexing upserts in place.
		67	vstore.add_documents(batch, ids=[c.id for c in batch])
95	print(f"✅ Batch {batch_num}/{total_batches} stored", flush=True)	68	print(f"✅ Batch {batch_num}/{total_batches} stored", flush=True)
96		69
97	return len(docs), len(all_chunks)	70	return len(docs), len(chunks)

10 rag/__init__.py

src/business/rag/__init__.py · 107 → 111 lines · 18 rows changed

RAG entry points called by the API. Adapts to the graph's return value and to Documents.

cortex-core · hand-written

Detects the fallback path by checking whether the top chunk has a rerank_score attribute.
Reads chunk.text and score attributes from custom dataclasses.

cortex-langgraph · LangChain / LangGraph

The graph returns a confidence value, so metrics use it directly (confidence == "high").
Reads chunk.page_content; scores live in metadata and are stripped before returning sources.

Why it matters: Metrics read an explicit field instead of inferring the branch from an object's attributes.

Document.metadata

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/__init__.py		cortex-langgraph src/business/rag/__init__.py
⋯ 18 unchanged lines hidden
1	"""	1	"""
2	RAG pipeline business logic entry-points.	2	RAG pipeline business logic entry-points.
3		3
4	query_rag → called by RAGController.query	4	query_rag → called by RAGController.query
5	ingest_pdfs → called by RAGController.upload (after file is saved to uploads dir)	5	ingest_pdfs → called by RAGController.upload (after file is saved to uploads dir)
6	"""	6	"""
7		7
8	from __future__ import annotations	8	from __future__ import annotations
9		9
10	from pathlib import Path	10	from pathlib import Path
11	from typing import Dict, List	11	from typing import Dict, List
12		12
13	from dotenv import load_dotenv	13	from dotenv import load_dotenv
14		14
15	from src.api.metrics import (	15	from src.api.metrics import (
16	RAG_QUERIES_TOTAL,	16	RAG_QUERIES_TOTAL,
17	RAG_RETRIEVAL_LOW_CONFIDENCE_TOTAL,	17	RAG_RETRIEVAL_LOW_CONFIDENCE_TOTAL,
18	RAG_RETRIEVAL_TOP_SCORE,	18	RAG_RETRIEVAL_TOP_SCORE,
19	)	19	)
20	from src.business.rag.retrieval import RAGPipeline	20	from src.business.rag.retrieval import RAGPipeline
21	from src.business.rag.index_builder import build_index	21	from src.business.rag.index_builder import build_index
22	from src.business.rag.vector_store import create_vector_store	22	from src.business.rag.vector_store import reset_vector_store
23	from src.utils.config import load_config	23	from src.utils.config import load_config
24		24
25	load_dotenv()	25	load_dotenv()
⋯ 3 unchanged lines hidden
26		26
27	# Docling extracts table structure locally (no external API/key needed), so	27	# Docling extracts table structure locally (no external API/key needed), so
28	# table extraction is always available.	28	# table extraction is always available.
29	TABLE_OCR_ENABLED: bool = True	29	TABLE_OCR_ENABLED: bool = True
30		30
31	_PROJECT_ROOT = Path(__file__).resolve().parents[3]	31	_PROJECT_ROOT = Path(__file__).resolve().parents[3]
		32
		33	# Metadata keys the graph adds for its own bookkeeping; not part of the
		34	# chunk's provenance, so they're not echoed back in `sources[].metadata`.
		35	_SCORE_KEYS = ("vector_score", "rerank_score")
32		36
33		37
34	def rag_persist_dir() -> Path:	38	def rag_persist_dir() -> Path:
⋯ 21 unchanged lines hidden
35	"""Resolve the RAG vector-store location. Public (no leading underscore)	39	"""Resolve the RAG vector-store location. Public (no leading underscore)
36	because scripts/index_cli.py imports it too — the CLI and the API must	40	because scripts/index_cli.py imports it too — the CLI and the API must
37	always agree on where the index lives, or CLI-indexed content becomes	41	always agree on where the index lives, or CLI-indexed content becomes
38	invisible to /api/v1/rag/query and vice versa."""	42	invisible to /api/v1/rag/query and vice versa."""
39	cfg = load_config()	43	cfg = load_config()
40	dirs = cfg.get("directories", {})	44	dirs = cfg.get("directories", {})
41	path = _PROJECT_ROOT / dirs.get("rag_vectorstore_dir", "data/rag_vectorstore")	45	path = _PROJECT_ROOT / dirs.get("rag_vectorstore_dir", "data/rag_vectorstore")
42	path.mkdir(parents=True, exist_ok=True)	46	path.mkdir(parents=True, exist_ok=True)
43	return path	47	return path
44		48
45		49
46	def _rag_uploads_dir() -> Path:	50	def _rag_uploads_dir() -> Path:
47	cfg = load_config()	51	cfg = load_config()
48	dirs = cfg.get("directories", {})	52	dirs = cfg.get("directories", {})
49	path = _PROJECT_ROOT / dirs.get("rag_uploads_dir", "data/rag_uploads")	53	path = _PROJECT_ROOT / dirs.get("rag_uploads_dir", "data/rag_uploads")
50	path.mkdir(parents=True, exist_ok=True)	54	path.mkdir(parents=True, exist_ok=True)
51	return path	55	return path
52		56
53		57
54	def _make_pipeline() -> RAGPipeline:	58	def _make_pipeline() -> RAGPipeline:
55	return RAGPipeline(persist_dir=str(rag_persist_dir()))	59	return RAGPipeline(persist_dir=str(rag_persist_dir()))
56		60
57		61
58	async def query_rag(question: str) -> Dict:	62	async def query_rag(question: str) -> Dict:
59	"""Run the full RAG pipeline: retrieve → rerank → generate."""	63	"""Run the RAG graph: retrieve → rerank → (gate) → generate."""
60	pipeline = _make_pipeline()	64	pipeline = _make_pipeline()
61	answer_text, selected_chunks = pipeline.answer(question)	65	answer_text, selected_chunks, confidence = pipeline.answer(question)
62		66
63	RAG_QUERIES_TOTAL.inc()	67	RAG_QUERIES_TOTAL.inc()
64	# The top chunk only has rerank_score when the re-ranker found something	68	# The graph reports which branch ran. "high" means the re-ranker found
65	# above its relevance gate; otherwise select_context() fell back to raw	69	# something above its relevance gate; anything else means it fell back to
66	# vector order (plain RetrievedChunk, no rerank_score) — that fallback is	70	# raw vector order (or refused) — itself the strongest signal that
67	# itself the strongest signal that retrieval quality is degrading.	71	# retrieval quality is degrading.
68	if selected_chunks and hasattr(selected_chunks[0], "rerank_score"):	72	if confidence == "high" and selected_chunks:
69	RAG_RETRIEVAL_TOP_SCORE.observe(selected_chunks[0].rerank_score)	73	RAG_RETRIEVAL_TOP_SCORE.observe(selected_chunks[0].metadata["rerank_score"])
70	else:	74	else:
71	RAG_RETRIEVAL_LOW_CONFIDENCE_TOTAL.inc()	75	RAG_RETRIEVAL_LOW_CONFIDENCE_TOTAL.inc()
72		76
73	sources: List[Dict] = [	77	sources: List[Dict] = [
74	{	78	{
75	"text": chunk.text[:400],	79	"text": chunk.page_content[:400],
76	"metadata": chunk.metadata,	80	"metadata": {k: v for k, v in chunk.metadata.items() if k not in _SCORE_KEYS},
77	# ReRankedChunk has rerank_score; fallback RetrievedChunk only has vector_score	81	# Re-ranked chunks carry rerank_score; fallback chunks only vector_score.
78	"score": round(float(getattr(chunk, "rerank_score", chunk.vector_score)), 4),	82	"score": round(float(chunk.metadata.get("rerank_score", chunk.metadata.get("vector_score", 0.0))), 4),
79	}	83	}
80	for chunk in selected_chunks	84	for chunk in selected_chunks
81	]	85	]
⋯ 11 unchanged lines hidden
82		86
83	return {"answer": answer_text, "sources": sources}	87	return {"answer": answer_text, "sources": sources}
84		88
85		89
86	async def ingest_pdfs(uploads_dir: Path) -> Dict:	90	async def ingest_pdfs(uploads_dir: Path) -> Dict:
87	"""	91	"""
88	Build / rebuild the vector index from all PDFs in uploads_dir.	92	Build / rebuild the vector index from all PDFs in uploads_dir.
89	The collection is reset first so stale chunks from previous uploads	93	The collection is reset first so stale chunks from previous uploads
90	are removed before the new PDF is indexed.	94	are removed before the new PDF is indexed.
91	"""	95	"""
92	persist_dir = rag_persist_dir()	96	persist_dir = rag_persist_dir()
93		97
94	# Wipe old chunks — upsert never deletes, so stale content from previous	98	# Wipe old chunks — upsert never deletes, so stale content from previous
95	# uploads would otherwise persist and pollute query results.	99	# uploads would otherwise persist and pollute query results.
96	create_vector_store(persist_dir=str(persist_dir)).reset()	100	reset_vector_store(persist_dir=str(persist_dir))
97		101
98	docs_count, chunks_count = build_index(	102	docs_count, chunks_count = build_index(
99	data_dir=uploads_dir,	103	data_dir=uploads_dir,
⋯ 8 unchanged lines hidden
100	persist_dir=persist_dir,	104	persist_dir=persist_dir,
101	max_context_chars=500_000, # large enough to capture entire papers including tables	105	max_context_chars=500_000, # large enough to capture entire papers including tables
102	)	106	)
103	return {	107	return {
104	"docs_indexed": docs_count,	108	"docs_indexed": docs_count,
105	"chunks_indexed": chunks_count,	109	"chunks_indexed": chunks_count,
106	"table_ocr_enabled": TABLE_OCR_ENABLED,	110	"table_ocr_enabled": TABLE_OCR_ENABLED,
107	}	111	}

11 pdf_digest.py

src/business/rag/pdfingest/pdf_digest.py · 150 → 122 lines · 159 rows changed

Reading PDFs. Docling does the parsing on both sides (locally, no API key); only the wrapper changes.

cortex-core · hand-written

Walks Docling's document model by hand: text blocks, then tables appended after a TABLES FROM DOCUMENT IMAGES separator.
Returns a custom IngestedDocument dataclass.

cortex-langgraph · LangChain / LangGraph

langchain_docling.DoclingLoader turns each PDF into LangChain Documents.
Tables are exported as Markdown inline, where they appear, next to the text that explains them.
Noise (page headers, footers, captions) is dropped by Docling's own labels= filter; a MetaExtractor adds document metadata.

Why it matters: Keeping a table next to its explaining text gives the retriever better chunks.

DoclingLoaderDocument

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/pdfingest/pdf_digest.py		cortex-langgraph src/business/rag/pdfingest/pdf_digest.py
1	"""Reusable PDF ingestion pipeline for RAG, powered by Docling.	1	"""PDF ingestion for RAG — LangChain `DoclingLoader` edition.
2		2
3	Docling (https://github.com/docling-project/docling) parses layout, OCR, and	3	Docling (https://github.com/docling-project/docling) still does the actual
4	table structure fully locally — no external API, no API key, and no per-page	4	work — layout analysis, OCR (RapidOCR), and table structure — fully locally,
5	cost, unlike the Unstructured-API-based table extraction this module replaces.	5	no API key. What changed is the wrapper: instead of walking Docling's object
		6	model by hand, `langchain_docling.DoclingLoader` turns each PDF into a
		7	LangChain `Document`, the unit every other LangChain component (splitters,
		8	vector stores, retrievers) consumes.
6		9
7	Converts PDFs in a directory into structured artifacts:	10	Differences from cortex-core, all deliberate:
8	- extracted text blocks (section headers, narrative text, list items)	11	- Tables are exported as Markdown inline, where they appear in the document,
9	- tables, extracted directly as Markdown (no image round-trip needed)	12	instead of being appended after a "TABLES FROM DOCUMENT IMAGES" separator.
10	- combined text trimmed to a configurable character budget	13	Each table now sits next to the text that explains it.
11		14	- Cleaning is done by Docling's own export filter (`labels=`): only titles,
12	Downstream steps (chunk → embed → index) can consume the returned	15	section headers, body text, list items and tables are kept — page headers,
13	``IngestedDocument`` objects without dealing with side effects or prints.	16	footers, footnotes and captions are dropped as noise, same as before.
		17	- Document-level metadata is added by a custom `MetaExtractor`
		18	(source_id, page/table counts) — LangChain's hook for metadata enrichment.
14	"""	19	"""
15		20
16	from __future__ import annotations	21	from __future__ import annotations
17		22
18	from dataclasses import dataclass
19	from pathlib import Path	23	from pathlib import Path
20	from typing import Dict, List, Optional	24	from typing import Any, Dict, List
21		25
22	from docling.datamodel.base_models import InputFormat	26	from docling.datamodel.base_models import InputFormat
23	from docling.datamodel.pipeline_options import PdfPipelineOptions	27	from docling.datamodel.pipeline_options import PdfPipelineOptions
24	from docling.document_converter import DocumentConverter, PdfFormatOption	28	from docling.document_converter import DocumentConverter, PdfFormatOption
		29	from docling_core.types.doc import DocItemLabel, DoclingDocument
		30	from langchain_core.documents import Document
		31	from langchain_docling import DoclingLoader
		32	from langchain_docling.loader import BaseMetaExtractor, ExportType, MetaExtractor
25		33
26		34
27	# Docling's DocItemLabel values (as plain strings) that correspond to the	35	# Docling item labels kept in the exported text. Page headers/footers,
28	# "Title / NarrativeText / ListItem" categories the rest of the pipeline	36	# footnotes, and captions are dropped as noise (cleaning step).
29	# expects. Page headers/footers, footnotes, and captions are dropped as noise.	37	USEFUL_LABELS = {
30	USEFUL_LABELS = {"section_header", "title", "text", "list_item"}	38	DocItemLabel.TITLE,
		39	DocItemLabel.SECTION_HEADER,
		40	DocItemLabel.TEXT,
		41	DocItemLabel.PARAGRAPH,
		42	DocItemLabel.LIST_ITEM,
		43	}
31		44
32		45
33	@dataclass	46	class CortexMetaExtractor(MetaExtractor):
34	class IngestedDocument:	47	"""Document-level metadata enrichment.
35	source_id: str	48
36	text_blocks: List[str]	49	`source_id` (the file name) is the key the rest of the pipeline uses for
37	table_text: str	50	provenance and stable chunk ids — the same key cortex-core used.
38	combined_text: str	51	"""
39	metadata: Dict[str, int]	52
		53	def extract_dl_doc_meta(self, file_path: str, dl_doc: DoclingDocument) -> Dict[str, Any]:
		54	return {
		55	"source": str(file_path),
		56	"source_id": Path(file_path).name,
		57	"page_count": len(dl_doc.pages),
		58	"table_count": len(dl_doc.tables),
		59	}
40		60
41		61
42	def _build_converter(do_ocr: bool, do_table_structure: bool) -> DocumentConverter:	62	def _build_converter(do_ocr: bool, do_table_structure: bool) -> DocumentConverter:
⋯ 3 unchanged lines hidden
43	pipeline_options = PdfPipelineOptions(do_ocr=do_ocr, do_table_structure=do_table_structure)	63	pipeline_options = PdfPipelineOptions(do_ocr=do_ocr, do_table_structure=do_table_structure)
44	return DocumentConverter(	64	return DocumentConverter(
45	format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)}	65	format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)}
46	)	66	)
47		67
48		68
49	def _load_docling_document(pdf_path: Path, strategy: str = "hi_res", extract_tables: bool = True):	69	def build_pdf_loader(
50	"""Convert a PDF with Docling. strategy: 'hi_res' (OCR on, slower, better for scans) \| 'fast' (no OCR)."""	70	pdf_paths: List[Path],
51	converter = _build_converter(do_ocr=(strategy == "hi_res"), do_table_structure=extract_tables)	71	include_table_images: bool = True,
52	return converter.convert(str(pdf_path)).document	72	pdf_strategy: str = "hi_res",
		73	meta_extractor: BaseMetaExtractor \| None = None,
		74	) -> DoclingLoader:
		75	"""Configure a DoclingLoader for the given PDFs.
		76
		77	pdf_strategy: 'hi_res' (OCR on, slower, needed for scanned pages) \| 'fast' (no OCR).
		78	include_table_images: whether to run Docling's table structure recognition
		79	and keep tables in the exported text.
		80	"""
		81	labels = set(USEFUL_LABELS)
		82	if include_table_images:
		83	labels.add(DocItemLabel.TABLE)
		84
		85	return DoclingLoader(
		86	file_path=[str(p) for p in pdf_paths],
		87	converter=_build_converter(do_ocr=(pdf_strategy == "hi_res"), do_table_structure=include_table_images),
		88	# MARKDOWN = one Document per PDF. Chunking is a separate, LangChain
		89	# step (see chunk.py) so its size/overlap stay under our control.
		90	export_type=ExportType.MARKDOWN,
		91	md_export_kwargs={"labels": labels, "image_placeholder": ""},
		92	meta_extractor=meta_extractor or CortexMetaExtractor(),
		93	)
53		94
54		95
55	def _text_blocks_from_document(doc) -> List[str]:	96	def load_pdf_documents(
56	return [	97	data_dir: Path,
57	item.text.strip()
58	for item in doc.texts
59	if str(item.label) in USEFUL_LABELS and getattr(item, "text", None) and item.text.strip()
60	]
61
62
63	def _table_text_from_document(doc) -> str:
64	table_texts = [table.export_to_markdown(doc) for table in doc.tables]
65	return "\n\n".join(t for t in table_texts if t.strip())
66
67
68	def ingest_single_pdf(
69	pdf_path: Path,
70	image_output_dir: Optional[Path] = None,
71	max_context_chars: int = 12_000,	98	max_context_chars: int = 12_000,
72	include_table_images: bool = True,	99	include_table_images: bool = True,
73	pdf_strategy: str = "hi_res",	100	pdf_strategy: str = "hi_res",
74	) -> IngestedDocument:	101	) -> List[Document]:
75	"""Ingest one PDF and return structured text artifacts.	102	"""Load every PDF under data_dir (recursively) as one LangChain Document each.
76		103
77	pdf_strategy: 'hi_res' (OCR on, slower, needed for scanned pages) \| 'fast' (no OCR).	104	Each document's text is trimmed to max_context_chars, the same budget
78	include_table_images: whether to run Docling's table structure recognition.	105	cortex-core applied.
79	image_output_dir: unused, kept for backward compatibility with older callers.
80	"""	106	"""
81	doc = _load_docling_document(pdf_path, strategy=pdf_strategy, extract_tables=include_table_images)	107	pdf_files = sorted(Path(data_dir).rglob("*.pdf"))
82
83	text_blocks = _text_blocks_from_document(doc)
84	table_text = _table_text_from_document(doc) if include_table_images else ""
85
86	pdf_text_only = "\n\n".join(text_blocks)
87
88	combined_text = pdf_text_only
89	separator = (
90	"\n\n" + "=" * 80 + "\n" + "TABLES FROM DOCUMENT IMAGES:\n" + "=" * 80 + "\n"
91	)
92
93	if table_text:
94	combined_text = pdf_text_only + separator + table_text
95
96	# Trim to budget. Table text is always preserved in full; body text is trimmed
97	# to fill whatever budget remains. This prevents tables from being silently
98	# dropped when the combined document exceeds max_context_chars.
99	if len(combined_text) > max_context_chars:
100	if table_text:
101	table_block = separator + table_text
102	available_for_pdf = max_context_chars - len(table_block)
103	if available_for_pdf > 0:
104	combined_text = pdf_text_only[:available_for_pdf] + table_block
105	else:
106	# table text alone already exceeds budget — keep it, trim body
107	combined_text = table_text[:max_context_chars]
108	else:
109	combined_text = pdf_text_only[:max_context_chars]
110
111	metadata = {
112	"text_block_count": len(text_blocks),
113	"table_count": len(doc.tables),
114	"total_chars": len(combined_text),
115	}
116
117	return IngestedDocument(
118	source_id=pdf_path.name,
119	text_blocks=text_blocks,
120	table_text=table_text,
121	combined_text=combined_text,
122	metadata=metadata,
123	)
124
125
126	def ingest_directory(
127	data_dir: Path,
128	image_output_dir: Optional[Path] = None,
129	max_context_chars: int = 12_000,
130	include_table_images: bool = True,
131	pdf_strategy: str = "hi_res",
132	) -> List[IngestedDocument]:
133	"""Ingest all PDFs in a directory and its subdirectories, and return
134	structured results. pdf_strategy: 'hi_res' (OCR on, slower) \| 'fast' (no OCR).
135	"""
136	pdf_files = list(Path(data_dir).rglob("*.pdf"))
137	if not pdf_files:	108	if not pdf_files:
138	raise FileNotFoundError(f"No PDF files found in {data_dir} (searched recursively)")	109	raise FileNotFoundError(f"No PDF files found in {data_dir} (searched recursively)")
139		110
140	results: List[IngestedDocument] = []	111	loader = build_pdf_loader(
141	for pdf_path in pdf_files:	112	pdf_files,
142	results.append(	113	include_table_images=include_table_images,
143	ingest_single_pdf(	114	pdf_strategy=pdf_strategy,
144	pdf_path=pdf_path,	115	)
145	max_context_chars=max_context_chars,	116	documents: List[Document] = []
146	include_table_images=include_table_images,	117	for doc in loader.lazy_load():
147	pdf_strategy=pdf_strategy,	118	text = doc.page_content.strip()[:max_context_chars]
148	)	119	doc.page_content = text
149	)	120	doc.metadata["total_chars"] = len(text)
150	return results	121	documents.append(doc)
		122	return documents

12 chunk.py

src/business/rag/pdfingest/chunk.py · 71 → 85 lines · 121 rows changed

Splitting documents into chunks. Same size and overlap; smarter cut points.

cortex-core · hand-written

A fixed character window: every chunk is exactly chunk_size characters, cut wherever the window ends, often mid-word.

cortex-langgraph · LangChain / LangGraph

RecursiveCharacterTextSplitter with the same defaults (800 characters, 100 overlap) tries paragraph breaks first, then lines, sentences and words.
Keeps cortex-core's stable chunk ids (a hash of source, offset and text), so re-indexing overwrites instead of duplicating.

Why it matters: Chunks that end on natural boundaries read better to both the embedding model and the LLM.

RecursiveCharacterTextSplitter

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/pdfingest/chunk.py		cortex-langgraph src/business/rag/pdfingest/chunk.py
1	# The main chunking logic is implemented here.	1	"""Chunking — LangChain `RecursiveCharacterTextSplitter` edition.
2		2
3	from dataclasses import dataclass	3	cortex-core used a fixed character window: every chunk was exactly
4	from importlib import metadata	4	`chunk_size` characters, cut wherever the window ended — often mid-word or
5	from typing import List, Dict, Any	5	mid-sentence. `RecursiveCharacterTextSplitter` has the same size/overlap knobs
		6	but tries to cut on paragraph breaks first, then lines, then sentences, then
		7	words, so chunks end on natural boundaries.
		8
		9	Kept from cortex-core:
		10	- Same defaults (800 chars, 100 overlap).
		11	- Stable, content-derived chunk ids (sha1 of source + offset + text prefix),
		12	so re-indexing the same PDF overwrites chunks instead of duplicating them.
		13	- Chunk metadata: source_id, chunk_start/chunk_end offsets, chunk_strategy,
		14	and a `section` tag ("table" vs "text").
		15	"""
		16
		17	from __future__ import annotations
		18
6	import hashlib	19	import hashlib
		20	from typing import Dict, List
7		21
8	@dataclass	22	from langchain_core.documents import Document
9	class Chunk:	23	from langchain_text_splitters import RecursiveCharacterTextSplitter
10	chunk_id: str
11	text: str
12	metadata: Dict[str, Any]
13		24
14	class Chunker:	25
		26	def build_text_splitter(chunk_size: int = 800, overlap: int = 100) -> RecursiveCharacterTextSplitter:
		27	assert overlap < chunk_size, "overlap must be smaller than chunk_size"
		28	return RecursiveCharacterTextSplitter(
		29	chunk_size=chunk_size,
		30	chunk_overlap=overlap,
		31	# Records each chunk's character offset in the source text as
		32	# metadata["start_index"] — what cortex-core called chunk_start.
		33	add_start_index=True,
		34	)
		35
		36
		37	def split_documents(
		38	documents: List[Document],
		39	chunk_size: int = 800,
		40	overlap: int = 100,
		41	) -> List[Document]:
		42	"""Split loaded PDF documents into chunks with stable ids and provenance metadata.
		43
		44	Each returned Document has `id` set, so vector stores upsert by it.
15	"""	45	"""
16	Deterministic text chunker with overlap.	46	splitter = build_text_splitter(chunk_size=chunk_size, overlap=overlap)
		47	strategy_name = f"recursive_{chunk_size}_overlap_{overlap}"
17		48
18	Single source of truth for chunking.	49	chunks = splitter.split_documents(documents)
		50	for chunk in chunks:
		51	start = int(chunk.metadata.pop("start_index", 0))
		52	chunk.metadata.update({
		53	"chunk_start": start,
		54	"chunk_end": start + len(chunk.page_content),
		55	"chunk_strategy": strategy_name,
		56	"section": "table" if is_table_chunk(chunk.page_content) else "text",
		57	})
		58	# Vector stores only accept flat, scalar metadata values.
		59	chunk.metadata = {k: v for k, v in chunk.metadata.items() if isinstance(v, (str, int, float, bool))}
		60	chunk.id = build_chunk_id(chunk.page_content, start, chunk.metadata)
		61	return chunks
		62
		63
		64	def is_table_chunk(text: str) -> bool:
		65	"""A chunk is tagged "table" when most of its lines are Markdown table rows.
		66
		67	Docling exports tables inline as Markdown (`\| a \| b \|`), so the rows are
		68	easy to recognise. cortex-core could tag by offset instead because it
		69	appended every table after a fixed separator.
19	"""	70	"""
		71	lines = [line.strip() for line in text.splitlines() if line.strip()]
		72	if not lines:
		73	return False
		74	table_rows = sum(1 for line in lines if line.startswith("\|") and line.endswith("\|"))
		75	return table_rows / len(lines) >= 0.5
20		76
21	def __init__(
22	self,
23	chunk_size: int = 800,
24	overlap: int = 100,
25	strategy_name: str = "char_window_800_overlap_100",
26	):
27	assert overlap < chunk_size, "overlap must be smaller than chunk_size"
28		77
29	self.chunk_size = chunk_size	78	def build_chunk_id(chunk_text: str, start: int, metadata: Dict) -> str:
30	self.overlap = overlap	79	"""Stable chunk id: same PDF + same offset + same text → same id."""
31	self.strategy_name = strategy_name	80	base = (
32		81	str(metadata.get("source_id", "")) # PDF file name
33	def split(self,text: str, source_metadata: Dict,	82	+ str(start)
34	) -> List[Chunk]:	83	+ chunk_text[:100]
35	"""	84	)
36	Split raw text into overlapping chunks.	85	return hashlib.sha1(base.encode("utf-8")).hexdigest()
37	Args:
38	text: full document text
39	source_metadata: metadata from ingestion (pdf path, page map, etc.)
40	Returns:
41	List[Chunk]
42	"""
43	if not text or not text.strip():
44	return []
45
46	chunks: List[Chunk] = []
47	start = 0
48	text_length = len(text)
49
50	while start < text_length:
51	end = min(start + self.chunk_size, text_length)
52	chunk_text = text[start:end].strip()
53
54	if chunk_text:
55	chunk_id = self._build_chunk_id(chunk_text=chunk_text,start=start,source_metadata=source_metadata,)
56	metadata = {**source_metadata,"chunk_start": start,"chunk_end": end,"chunk_strategy": self.strategy_name,}
57	chunks.append(Chunk(chunk_id=chunk_id,text=chunk_text,metadata=metadata,))
58	start = end - self.overlap
59	if end >= text_length:
60	break
61	return chunks
62	def _build_chunk_id(self,chunk_text: str,start: int,source_metadata: Dict,) -> str:
63	"""
64	Build stable chunk ID
65	"""
66	base = (
67	source_metadata.get("source_id", "") #Pdf files Name
68	+ str(start)
69	+ chunk_text[:100]
70	)
71	return hashlib.sha1(base.encode("utf-8")).hexdigest()

13 re_ranker.py

src/business/rag/re_ranker/re_ranker.py · 112 → 79 lines · 157 rows changed

Re-ranking: score each candidate, drop weak ones, keep the best. The one part with no LangChain equivalent, so the policy stays hand-written.

cortex-core · hand-written

A ReRanker class: score → gate on a minimum score → sort → keep the top n, on custom chunk dataclasses.

cortex-langgraph · LangChain / LangGraph

A rerank_documents() function on LangChain Documents; scores ride in metadata (vector_score, rerank_score).
The scorer is any LangChain BaseCrossEncoder.
The "everything was rejected" decision moved out into the RAG graph.

Why it matters: LangChain's CrossEncoderReranker discards scores and has no threshold, but the threshold is the RAG path's first defence against hallucination. So the policy stays custom.

BaseCrossEncoderDocument.metadata

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/re_ranker/re_ranker.py		cortex-langgraph src/business/rag/re_ranker/re_ranker.py
1	# Some part of this project is for the RAG system and some parts are for the chatbot system.	1	"""Re-ranking logic: score → gate → sort → truncate.
2	# The Re-Ranker is only used in the RAG system.	2
		3	This is the part of cortex-core's re-ranker that has no LangChain equivalent.
		4	LangChain's `CrossEncoderReranker` (a document compressor) keeps the top-n by
		5	score but throws the scores away and has no relevance threshold — and the
		6	threshold is the RAG path's first hallucination firewall. So the scoring model
		7	is a LangChain `BaseCrossEncoder`, and the policy around it stays ours.
		8
		9	What moved out: cortex-core's `orchestrator.select_context()` (what to do when
		10	the gate rejects everything). That decision is now a conditional edge in the
		11	RAG graph (retrieval.py), where it's visible as a branch.
		12
		13	Chunks are LangChain `Document`s; scores ride along in their metadata:
		14	metadata["vector_score"] — distance from the vector store (lower = closer)
		15	metadata["rerank_score"] — cross-encoder relevance (higher = more relevant)
		16	"""
		17
		18	from __future__ import annotations
3		19
4	from typing import List	20	from typing import List
5		21
		22	from langchain_core.cross_encoders import BaseCrossEncoder
		23	from langchain_core.documents import Document
		24
6	from .config import ReRankerConfig	25	from .config import ReRankerConfig
7	from .interface import (
8	RetrievedChunk,
9	ReRankedChunk,
10	ReRankScorer,
11	)
12		26
13		27
14	class ReRanker:	28	def rerank_documents(
		29	query: str,
		30	documents: List[Document],
		31	scorer: BaseCrossEncoder,
		32	config: ReRankerConfig,
		33	) -> List[Document]:
15	"""	34	"""
16	Re-ranking orchestrator.	35	Execute the re-ranking pipeline.
17		36
18	Responsibilities:	37	Returns re-ranked, gated, sorted, truncated chunks — possibly an empty
19	- Call the scoring model	38	list when nothing clears the relevance gate.
20	- Apply gating logic	39	"""
21	- Sort chunks by relevance	40	if not documents:
22	- Truncate to top-n	41	return []
23	- Return prompt-ready chunks
24		42
25	This class contains NO ML logic.	43	# 1. Limit input size (recall control)
		44	candidates = documents[: config.top_k_input]
		45
		46	# 2. Score relevance (model-dependent)
		47	scores = scorer.score([(query, doc.page_content) for doc in candidates])
		48	scored = [
		49	Document(
		50	id=doc.id,
		51	page_content=doc.page_content,
		52	metadata={**doc.metadata, "rerank_score": float(score)},
		53	)
		54	for doc, score in zip(candidates, scores)
		55	]
		56
		57	# 3. Apply gating (hallucination control)
		58	gated = [doc for doc in scored if doc.metadata["rerank_score"] >= config.min_score]
		59	if not gated:
		60	return []
		61
		62	# 4. Sort by relevance score (descending)
		63	ranked = sorted(gated, key=lambda doc: final_score(doc, config), reverse=True)
		64
		65	# 5. Truncate to top-n (context control)
		66	return ranked[: config.top_n_output]
		67
		68
		69	def final_score(doc: Document, config: ReRankerConfig) -> float:
26	"""	70	"""
		71	Compute final ranking score.
		72	Supports optional blending with vector score.
		73	"""
		74	rerank_score = doc.metadata["rerank_score"]
		75	if not config.blend_with_vector_score:
		76	return rerank_score
27		77
28	def __init__(	78	alpha = config.blend_alpha
29	self,	79	return alpha * rerank_score + (1 - alpha) * doc.metadata.get("vector_score", 0.0)
30	scorer: ReRankScorer,
31	config: ReRankerConfig,
32	) -> None:
33	self.scorer = scorer
34	self.config = config
35
36	def re_rank(
37	self,
38	query: str,
39	chunks: List[RetrievedChunk],
40	) -> List[ReRankedChunk]:
41	"""
42	Execute the re-ranking pipeline.
43
44	Args:
45	query: User query
46	chunks: Retrieved chunks from VectorDB
47
48	Returns:
49	List of re-ranked, filtered, sorted chunks
50	"""
51
52	if not chunks:
53	return []
54
55	# 1. Limit input size (recall control)
56	candidates = chunks[: self.config.top_k_input]
57
58	# 2. Score relevance (model-dependent)
59	scored_chunks = self.scorer.score(
60	query=query,
61	chunks=candidates,
62	)
63
64	# 3. Apply gating (hallucination control)
65	gated_chunks = self._apply_gating(scored_chunks)
66
67	if not gated_chunks:
68	return []
69
70	# 4. Sort by relevance score (descending)
71	ranked_chunks = sorted(
72	gated_chunks,
73	key=self._final_score,
74	reverse=True,
75	)
76
77	# 5. Truncate to top-n (context control)
78	return ranked_chunks[: self.config.top_n_output]
79
80	# =========================
81	# Internal helpers
82	# =========================
83
84	def _apply_gating(
85	self,
86	chunks: List[ReRankedChunk],
87	) -> List[ReRankedChunk]:
88	"""
89	Drop chunks below minimum relevance threshold.
90	"""
91	return [
92	chunk
93	for chunk in chunks
94	if chunk.rerank_score >= self.config.min_score
95	]
96
97	def _final_score(
98	self,
99	chunk: ReRankedChunk,
100	) -> float:
101	"""
102	Compute final ranking score.
103	Supports optional blending with vector score.
104	"""
105	if not self.config.blend_with_vector_score:
106	return chunk.rerank_score
107
108	alpha = self.config.blend_alpha
109	return (
110	alpha * chunk.rerank_score
111	+ (1 - alpha) * chunk.vector_score
112	)

14 cross_encoder.py

src/business/rag/re_ranker/cross_encoder.py · 110 → 60 lines · 120 rows changed

The model that scores how relevant each chunk is to the question.

cortex-core · hand-written

Drives the model by hand: tokenizer + AutoModelForSequenceClassification, manual batching, torch.no_grad().

cortex-langgraph · LangChain / LangGraph

sentence-transformers' CrossEncoder does the batching; a small class adapts it to LangChain's BaseCrossEncoder.score().
Keeps raw logits (activation_fn=Identity) so the existing min_score threshold gates exactly the same chunks.

Why it matters: Implementing the interface from langchain_core avoids depending on langchain-community, which is being retired.

BaseCrossEncodersentence-transformers

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/re_ranker/cross_encoder.py		cortex-langgraph src/business/rag/re_ranker/cross_encoder.py
1	"""	1	"""Cross-encoder scorer — implements LangChain's `BaseCrossEncoder` interface.
2	Here is the Business Logic of the Re-Ranker using Cross Encoder model.
3	VectorDB → RetrievedChunk[]
4	↓
5	CrossEncoderReRanker.score()
6	↓
7	ReRankedChunk[]
8	↓
9	ReRanker (gating, sorting, truncation)
10	↓
11	Prompt Builder
12		2
13	In a simple meaning, Cross Encoder is a sequence classification model	3	A cross-encoder is a sequence-classification model that reads the query and a
14	that takes two texts as input and outputs a single score indicating their relevance.	4	document together and outputs one relevance score. It is slower than vector
		5	similarity (one forward pass per pair) but much more accurate, which is why
		6	it runs only on the top-k candidates the vector store already found.
		7
		8	cortex-core drove the model by hand: tokenizer + AutoModelForSequenceClassification
		9	+ manual batching + torch.no_grad(). sentence-transformers' `CrossEncoder` does
		10	all of that; this class adapts it to LangChain's `BaseCrossEncoder`
		11	(`score(text_pairs) -> List[float]`) so it plugs into anything in the
		12	LangChain ecosystem that accepts a cross-encoder.
		13
		14	Why not langchain_community's HuggingFaceCrossEncoder: langchain-community is
		15	being retired. The interface lives in langchain_core, so implementing it here
		16	is ten lines.
		17
		18	Score scale — kept identical to cortex-core on purpose: sentence-transformers
		19	applies a sigmoid to single-label models by default, which would turn scores
		20	into 0-1 probabilities. activation_fn=Identity keeps the raw logits cortex-core
		21	produced, so ReRankerConfig.min_score gates exactly the same chunks in both
		22	projects (see the TestRerankerGate finding in tests/evals for why that
		23	threshold is on a logit scale).
15	"""	24	"""
16		25
17	from typing import List	26	from __future__ import annotations
		27
		28	from typing import List, Tuple
18		29
19	import torch	30	import torch
20	from transformers import AutoTokenizer, AutoModelForSequenceClassification	31	from langchain_core.cross_encoders import BaseCrossEncoder
21		32
22	from .interface import (
23	RetrievedChunk,
24	ReRankedChunk,
25	ReRankScorer,
26	)
27	from .config import ReRankerConfig	33	from .config import ReRankerConfig
28		34
29		35
30	class CrossEncoderReRanker(ReRankScorer):	36	class CrossEncoderScorer(BaseCrossEncoder):
31	"""	37	"""Scores (query, document) pairs with a Hugging Face cross-encoder model."""
32	Cross-encoder based re-ranker.
33
34	Uses a transformer model to jointly encode (query, document)
35	and output a relevance score.
36	"""
37		38
38	def __init__(self, config: ReRankerConfig) -> None:	39	def __init__(self, config: ReRankerConfig) -> None:
		40	from sentence_transformers import CrossEncoder
		41
39	self.config = config	42	self.config = config
40		43	device = config.device if torch.cuda.is_available() else "cpu"
41	self.device = torch.device(	44	self.model = CrossEncoder(
42	config.device if torch.cuda.is_available() else "cpu"	45	config.model_name,
		46	device=device,
		47	activation_fn=torch.nn.Identity(),
43	)	48	)
44		49
45	self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)	50	def score(self, text_pairs: List[Tuple[str, str]]) -> List[float]:
46	self.model = AutoModelForSequenceClassification.from_pretrained(	51	"""Relevance score for each (query, document) pair — raw logits."""
47	config.model_name	52	if not text_pairs:
48	).to(self.device)
49
50	self.model.eval()
51
52	def score(
53	self,
54	query: str,
55	chunks: List[RetrievedChunk],
56	) -> List[ReRankedChunk]:
57	"""
58	Compute relevance scores for query–chunk pairs.
59	"""
60
61	if not chunks:
62	return []	53	return []
63		54	scores = self.model.predict(
64	pairs = [(query, chunk.text) for chunk in chunks]	55	text_pairs,
65		56	batch_size=self.config.batch_size,
66	scores = self._batch_score(pairs)	57	show_progress_bar=False,
67		58	convert_to_numpy=True,
68	return [	59	)
69	ReRankedChunk(	60	return [float(s) for s in scores]
70	chunk_id=chunk.chunk_id,
71	text=chunk.text,
72	metadata=chunk.metadata,
73	vector_score=chunk.vector_score,
74	rerank_score=float(score),
75	)
76	for chunk, score in zip(chunks, scores)
77	]
78
79	# =========================
80	# Internal helpers
81	# =========================
82
83	def _batch_score(self, pairs: List[tuple[str, str]]) -> List[float]:
84	"""
85	Score query-document pairs in batches.
86	"""
87
88	all_scores: List[float] = []
89
90	with torch.no_grad():
91	for i in range(0, len(pairs), self.config.batch_size):
92	batch = pairs[i : i + self.config.batch_size]
93
94	inputs = self.tokenizer(
95	batch,
96	padding=True,
97	truncation=True,
98	return_tensors="pt",
99	).to(self.device)
100
101	outputs = self.model(**inputs)
102
103	logits = outputs.logits.squeeze(-1)
104
105	if logits.dim() == 0:
106	logits = logits.unsqueeze(0)
107
108	all_scores.extend(logits.cpu().tolist())
109
110	return all_scores

15 orchestrator.py

src/business/rag/re_ranker/orchestrator.py · 52 → 47 lines · 62 rows changed

What to do when the re-ranker rejects every chunk. Only in cortex-core; the right side is the graph code that replaced it.

cortex-core · hand-written

select_context() applies one of three policies: fail closed (refuse), fail open (use raw vector results) or hybrid (fall back and mark low confidence).

cortex-langgraph · LangChain / LangGraph

The same three policies are a conditional edge: _route_after_rerank() sends the graph to generate, fallback or refuse.

Why it matters: A policy that decides whether the LLM is called at all belongs in the visible flow, not in a helper.

add_conditional_edges

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/business/rag/re_ranker/orchestrator.py		cortex-langgraph src/business/rag/retrieval.py (excerpt: routing + graph)
1	"""	128	def _route_after_rerank(self, state: RAGState) -> str:
2	This file is after all the re-ranker code and tests, so the final vector results or chunks which are re-ranked and	129	"""The policy cortex-core kept in orchestrator.select_context()."""
3	filtered are passed to this orchestrator function to decide what to do next based on the selected policy and send	130	if state["context"]:
4	them to the LLM or not.	131	return "generate"
		132	if self.policy == "fail_closed":
		133	return "refuse"
		134	return "fallback"
5		135
6	Option A — Fail Closed (strict)	136	def _fallback(self, state: RAGState) -> RAGState:
7	If re-ranker returns empty → do not call LLM	137	# Fail open: raw vector order, flagged low-confidence so it shows up in metrics.
8	Return a safe refusal (“I don’t have enough information…”)	138	return {
9	Best for high-risk domains	139	"context": state["candidates"][: self.reranker_config.top_n_output],
10	Option B — Fail Open (fallback)	140	"confidence": "low",
11	If empty → fall back to top-k vector results	141	}
12	Proceed with caution
13	Best for exploratory/internal tools
14	Option C — Hybrid (recommended default)
15	Try strict re-ranking
16	If empty → fall back to vector results and mark low confidence
17	Log the event
18	"""
19	from typing import List, Tuple
20		142
21	from .re_ranker import ReRanker	143	def _refuse(self, state: RAGState) -> RAGState:
22	from .interface import RetrievedChunk, ReRankedChunk	144	return {"context": [], "confidence": "none", "answer": REFUSAL}
23		145
		146	def _generate(self, state: RAGState) -> RAGState:
		147	answer = self.generate_chain.invoke({
		148	"question": state["question"],
		149	"context": format_context(state["context"]),
		150	})
		151	return {"answer": answer.strip()}
24		152
25	def select_context(	153	# ------------------------------------------------------------------ graph
26	query: str,	154	def _build_graph(self):
27	retrieved_chunks: List[RetrievedChunk],	155	builder = StateGraph(RAGState)
28	reranker: ReRanker,	156	builder.add_node("retrieve", self._retrieve)
29	policy: str = "hybrid", # "fail_closed" \| "fail_open" \| "hybrid"	157	builder.add_node("rerank", self._rerank)
30	) -> Tuple[List[ReRankedChunk], str]:	158	builder.add_node("fallback", self._fallback)
31	"""	159	builder.add_node("refuse", self._refuse)
32	Select final context for the LLM after re-ranking.	160	builder.add_node("generate", self._generate)
33		161
34	Returns:	162	builder.add_edge(START, "retrieve")
35	chunks: final context chunks	163	builder.add_edge("retrieve", "rerank")
36	confidence: "high" \| "low" \| "none"	164	builder.add_conditional_edges(
37	"""	165	"rerank",
		166	self._route_after_rerank,
		167	{"generate": "generate", "fallback": "fallback", "refuse": "refuse"},
		168	)
		169	builder.add_edge("fallback", "generate")
		170	builder.add_edge("generate", END)
		171	builder.add_edge("refuse", END)
		172	return builder.compile()
38		173
39	reranked_chunks = reranker.re_rank(query, retrieved_chunks)	174	# -------------------------------------------------------------------- API
40
41	# Case 1: Re-ranker succeeded
42	if reranked_chunks:
43	return reranked_chunks, "high"
44
45	# Case 2: Fail-closed
46	if policy == "fail_closed":
47	return [], "none"
48
49	# Case 3: Fail-open / hybrid fallback
50	fallback_chunks = retrieved_chunks[: reranker.config.top_n_output]
51
52	return fallback_chunks, "low"

16 interface.py

src/business/rag/re_ranker/interface.py · 67 lines · only in cortex-core

Data contracts for the re-ranker. Only in cortex-core.

cortex-core · hand-written

RetrievedChunk and ReRankedChunk dataclasses, plus a ReRankScorer abstract class.

cortex-langgraph · LangChain / LangGraph

Not needed: every stage passes LangChain Documents, and the scorer interface is BaseCrossEncoder.

Why it matters: Using the framework's shared types means every component already agrees on the data shape.

DocumentBaseCrossEncoder

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-langgraph: No equivalent file. LangChain's Document (with scores in metadata) replaces the chunk dataclasses, and BaseCrossEncoder replaces the scorer interface.

cortex-core src/business/rag/re_ranker/interface.py
1	# Interface -> Contracts only
2	from dataclasses import dataclass
3	from typing import List, Dict, Any
4	from abc import ABC, abstractmethod
5
6
7	# =========================
8	# Data Contracts
9	# =========================
10
11	@dataclass
12	class RetrievedChunk:
13	"""
14	Chunk returned from the VectorDB before re-ranking.
15	"""
16	chunk_id: str
17	text: str
18	metadata: Dict[str, Any]
19	vector_score: float
20
21
22	@dataclass
23	class ReRankedChunk(RetrievedChunk):
24	"""
25	Chunk after re-ranking.
26	Inherits all fields from RetrievedChunk and adds rerank_score.
27	"""
28	rerank_score: float
29
30
31	# =========================
32	# Re-Ranker Interface
33	# =========================
34
35	class ReRankScorer(ABC):
36	"""
37	Abstract base class for all re-ranking implementations.
38
39	Responsibilities:
40	- Score relevance between a query and retrieved chunks
41	- Return chunks with rerank_score populated
42
43	Non-responsibilities:
44	- No gating (thresholding)
45	- No sorting or truncation
46	- No prompt construction
47	- No VectorDB access
48	"""
49
50	@abstractmethod
51	def score(
52	self,
53	query: str,
54	chunks: List[RetrievedChunk],
55	) -> List[ReRankedChunk]:
56	"""
57	Assign relevance scores to retrieved chunks.
58
59	Args:
60	query: User query string
61	chunks: List of RetrievedChunk from VectorDB
62
63	Returns:
64	List of ReRankedChunk with rerank_score populated.
65	Ordering is NOT enforced here.
66	"""
67	raise NotImplementedError

17 long_term_memory.py

src/memory/long_term_memory.py · 119 → 128 lines · 97 rows changed

The chatbot's long-term memory: remember, recall and forget, per user. Same behavior, different storage contract.

cortex-core · hand-written

Takes a custom vector-DB adapter, an embedder and a chunker, and embeds text itself before storing.

cortex-langgraph · LangChain / LangGraph

Takes any LangChain VectorStore (which embeds text itself) and an optional TextSplitter.
Stores memories as Documents with user and timestamp metadata.

Why it matters: One dependency instead of three, and the memory works with any LangChain vector store.

VectorStoreTextSplitterDocument

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/memory/long_term_memory.py		cortex-langgraph src/memory/long_term_memory.py
1	# No specific long‑term memory service is hard‑coded in long_term_memory.py.	1	"""Long-term (semantic) memory — LangChain edition.
2	# Replace it with Letta Long Term Memory or Langchain long term memory	2
3	from typing import List, Dict, TYPE_CHECKING	3	Same memory semantics as cortex-core (remember / recall / forget, scoped per
		4	user), but the storage is any LangChain `VectorStore` instead of a hand-written
		5	adapter, and chunking uses LangChain's `RecursiveCharacterTextSplitter`.
		6
		7	The vector store embeds text itself, so this class no longer needs an
		8	embedder of its own — that dependency moved into the store.
		9	"""
		10
4	from datetime import datetime	11	from datetime import datetime
		12	from typing import Dict, List, Optional
5		13
6	# Type checking only - avoids circular imports	14	from langchain_core.documents import Document
7	if TYPE_CHECKING:	15	from langchain_core.vectorstores import VectorStore
8	from src.memory.vectordb import ConversationVectorStoreBase	16	from langchain_text_splitters import RecursiveCharacterTextSplitter, TextSplitter
9		17
10		18
11	class LongTermMemory:	19	class LongTermMemory:
12	"""	20	"""
13	Cognitive long-term memory layer.	21	Cognitive long-term memory layer.
14	Owns memory semantics, not storage.	22	Owns memory semantics, not storage.
15		23	Uses dependency injection - the vector store (and optionally the text
16	Uses dependency injection - vectordb, embedder, and chunker are passed in,not imported.	24	splitter) are passed in, not imported. This makes the code more flexible
17	This makes the code more flexible and testable.	25	and testable.
18	"""	26	"""
19		27
20	def __init__(	28	def __init__(
21	self,	29	self,
22	vectordb: "ConversationVectorStoreBase", # Type hint (string to avoid circular import)	30	vectorstore: VectorStore,
23	embedder, # Object with embed() method	31	text_splitter: Optional[TextSplitter] = None,
24	chunker, # Object with split() method
25	):	32	):
26	"""	33	"""
27	Initialize long-term memory.
28
29	Args:	34	Args:
30	vectordb: Vector database instance (from src.memory.vectordb.VectorDB)	35	vectorstore: Any LangChain VectorStore that supports metadata
31	embedder: Embedding model (must have embed(text: str) or embed_query(text: str) method)	36	filtering on `user_id` (Chroma does).
32	chunker: Text chunker (must have split(text: str) method)	37	text_splitter: Used only by remember(); conversation turns are
		38	stored whole. Defaults to a 1000-char recursive splitter.
33	"""	39	"""
34	self.vectordb = vectordb	40	self.vectorstore = vectorstore
35	self.embedder = embedder	41	self.text_splitter = text_splitter or RecursiveCharacterTextSplitter(
36	self.chunker = chunker	42	chunk_size=1000, chunk_overlap=100
		43	)
37		44
38	def remember(	45	def remember(
39	self,	46	self,
⋯ 2 unchanged lines hidden
40	content: str,	47	content: str,
41	user_id: str,	48	user_id: str,
42	memory_type: str = "knowledge",	49	memory_type: str = "knowledge",
43	importance: int = 1,	50	importance: int = 1,
44	) -> None:	51	) -> None:
45
46	"""	52	"""
47	chunk the information, embed each chunk, store in vectordb with metadata	53	Store information in long-term memory: split it into chunks and store
48	Store information in long-term memory using a pluggable chunking strategy.	54	each chunk with user-scoped metadata (the store embeds them).
49	"""	55	"""
50	chunks = self.chunker.split(content)	56	chunks = self.text_splitter.split_text(content)
51		57	if not chunks:
52	for chunk in chunks:	58	return
53	embedding = self.embedder.embed(chunk)	59	docs = [
54		60	Document(page_content=chunk, metadata=self._metadata(user_id, memory_type, importance))
55	self.vectordb.add(	61	for chunk in chunks
56	ids=[self._build_id(user_id)],	62	]
57	embeddings=[embedding],	63	self.vectorstore.add_documents(docs, ids=[self._build_id(user_id, i) for i in range(len(docs))])
58	documents=[chunk],
59	metadatas=[{
60	"user_id": user_id,
61	"type": memory_type,
62	"importance": importance,
63	"created_at": datetime.utcnow().isoformat(),
64	}],
65	)
66		64
67	def recall(	65	def recall(
68	self,	66	self,
⋯ 3 unchanged lines hidden
69	query: str,	67	query: str,
70	user_id: str,	68	user_id: str,
71	top_k: int = 5,	69	top_k: int = 5,
72	) -> List[Dict]:	70	) -> List[Dict]:
73	"""	71	"""
74	Recall relevant memories for a user query.	72	Recall relevant memories for a user query.
		73
		74	Returns the same shape cortex-core did — [{"text", "metadata", "score"}] —
		75	so the prompt builder and the agent tool don't care which store is behind it.
		76	Score is the store's distance (lower = more similar).
75	"""	77	"""
76	embedding = self.embedder.embed(query)	78	results = self.vectorstore.similarity_search_with_score(
77		79	query,
78	return self.vectordb.search(	80	k=top_k,
79	embedding=embedding,	81	filter={"user_id": user_id},
80	top_k=top_k,
81	filters={"user_id": user_id},
82	)	82	)
		83	return [
		84	{"text": doc.page_content, "metadata": doc.metadata, "score": score}
		85	for doc, score in results
		86	]
83		87
84	def forget_user(self, user_id: str) -> None:	88	def forget_user(self, user_id: str) -> None:
85	"""	89	"""
86	Delete all memories for a user.	90	Delete all memories for a user.
		91
		92	LangChain's generic VectorStore.delete() only takes ids, so this looks
		93	the user's ids up first via Chroma's metadata filter (`get(where=...)`).
87	"""	94	"""
88	self.vectordb.delete(filters={"user_id": user_id})	95	existing = self.vectorstore.get(where={"user_id": user_id})
		96	ids = existing.get("ids", [])
		97	if ids:
		98	self.vectorstore.delete(ids=ids)
89		99
90	def remember_conversation(	100	def remember_conversation(
91	self,	101	self,
⋯ 3 unchanged lines hidden
92	user_message: str,	102	user_message: str,
93	assistant_response: str,	103	assistant_response: str,
94	user_id: str,	104	user_id: str,
95	) -> None:	105	) -> None:
96	"""	106	"""
97	Store a single conversation turn (user + assistant exchange) as one atomic unit.	107	Store a single conversation turn (user + assistant exchange) as one atomic unit.
98
99	Skips chunking intentionally — conversation pairs are already small,	108	Skips chunking intentionally — conversation pairs are already small,
100	and splitting them would break their semantic coherence.	109	and splitting them would break their semantic coherence.
101	"""	110	"""
102	text = f"user: {user_message}\nassistant: {assistant_response}"	111	text = f"user: {user_message}\nassistant: {assistant_response}"
103	embedding = self.embedder.embed(text)	112	doc = Document(page_content=text, metadata=self._metadata(user_id, "conversation", 1))
104		113	self.vectorstore.add_documents([doc], ids=[self._build_id(user_id)])
105	self.vectordb.add(
106	ids=[self._build_id(user_id)],
107	embeddings=[embedding],
108	documents=[text],
109	metadatas=[{
110	"user_id": user_id,
111	"type": "conversation",
112	"importance": 1,
113	"created_at": datetime.utcnow().isoformat(),
114	}],
115	)
116		114
117	@staticmethod	115	@staticmethod
118	def _build_id(user_id: str) -> str:	116	def _metadata(user_id: str, memory_type: str, importance: int) -> Dict:
119	return f"{user_id}-{datetime.utcnow().timestamp()}"	117	return {
		118	"user_id": user_id,
		119	"type": memory_type,
		120	"importance": importance,
		121	"created_at": datetime.utcnow().isoformat(),
		122	}
		123
		124	@staticmethod
		125	def _build_id(user_id: str, index: int = 0) -> str:
		126	# The chunk index keeps ids unique when remember() stores several
		127	# chunks within the same timestamp tick.
		128	return f"{user_id}-{datetime.utcnow().timestamp()}-{index}"

18 vectordb.py

src/memory/vectordb.py · 131 → 66 lines · 127 rows changed

The vector store behind conversation memory, kept separate from the PDF index on purpose.

cortex-core · hand-written

A hand-written ChromaVectorDB adapter (add / search / delete) over chromadb.Client.
In chromadb 1.x, Client(Settings(persist_directory=...)) is in-memory, so long-term memory was lost on every restart.

cortex-langgraph · LangChain / LangGraph

langchain_chroma.Chroma(persist_directory=...) implements the VectorStore interface and really persists to disk.
Its own provider switch (CHAT_VECTOR_STORE_PROVIDER), separate from the RAG store's.

Why it matters: Besides removing the adapter, the migration fixed a real bug: memories now survive restarts.

ChromaVectorStore

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core src/memory/vectordb.py		cortex-langgraph src/memory/vectordb.py
		1	"""Conversation-memory vector store — LangChain edition.
		2
		3	This is the chatbot's long-term-memory index — a separate index from the RAG
		4	vector store (src/business/rag/vector_store.py). Kept separate by design: RAG
		5	chunks (PDF content) and per-user conversation history serve different
		6	purposes and shouldn't share an index.
		7
		8	cortex-core wrapped chromadb by hand in a `ChromaVectorDB` adapter with its own
		9	`add/search/delete` interface. LangChain's `VectorStore` is that interface, and
		10	`langchain_chroma.Chroma` implements it — including embedding the text itself,
		11	so callers pass strings instead of pre-computed vectors.
		12
		13	Behavior note: cortex-core created `chromadb.Client(Settings(persist_directory=...))`,
		14	which in chromadb 1.x is an in-memory client — the directory was never written
		15	and long-term memory was lost on every restart. Chroma(persist_directory=...)
		16	here really persists, so memories now survive restarts.
		17	"""
		18
1	import os	19	import os
2	from abc import ABC, abstractmethod	20	from typing import Optional
3	from typing import List, Dict, Optional	21
4	from chromadb import Client
5	from chromadb.config import Settings	22	from chromadb.config import Settings
6		23	from langchain_core.embeddings import Embeddings
7	# This is the chatbot's conversation/long-term-memory vector store — a	24	from langchain_core.vectorstores import VectorStore
8	# separate index from the RAG vector store (src/business/rag/vector_store.py).
9	# Kept separate by design: RAG chunks (PDF content) and per-user conversation
10	# history serve different purposes and shouldn't share an index.
11
12
13	class ConversationVectorStoreBase(ABC):
14	"""Interface every conversation/long-term-memory vector store backend must implement."""
15
16	@abstractmethod
17	def add(
18	self,
19	ids: List[str],
20	embeddings: List[List[float]],
21	documents: List[str],
22	metadatas: List[Dict],
23	) -> None: ...
24
25	@abstractmethod
26	def search(
27	self,
28	embedding: List[float],
29	top_k: int = 5,
30	filters: Optional[Dict] = None,
31	) -> List[Dict]: ...
32
33	@abstractmethod
34	def delete(self, filters: Dict) -> None: ...
35
36
37	class ChromaVectorDB(ConversationVectorStoreBase):
38	"""
39	Low-level vector database adapter.
40	No AI logic. No memory semantics.
41	"""
42
43	def __init__(
44	self,
45	collection_name: str,
46	persist_directory: str = "./chroma",
47	):
48	self.client = Client(
49	Settings(persist_directory=persist_directory)
50	)
51	self.collection = self.client.get_or_create_collection(
52	name=collection_name
53	)
54
55	def add(
56	self,
57	ids: List[str],
58	embeddings: List[List[float]],
59	documents: List[str],
60	metadatas: List[Dict],
61	) -> None:
62	self.collection.add(
63	ids=ids,
64	embeddings=embeddings,
65	documents=documents,
66	metadatas=metadatas,
67	)
68
69	def search(
70	self,
71	embedding: List[float],
72	top_k: int = 5,
73	filters: Optional[Dict] = None,
74	) -> List[Dict]:
75	results = self.collection.query(
76	query_embeddings=[embedding],
77	n_results=top_k,
78	where=filters,
79	)
80
81	return [
82	{
83	"text": doc,
84	"metadata": meta,
85	"score": score,
86	}
87	for doc, meta, score in zip(
88	results["documents"][0],
89	results["metadatas"][0],
90	results["distances"][0],
91	)
92	]
93
94	def delete(self, filters: Dict) -> None:
95	self.collection.delete(where=filters)
96
97
98	# Backward-compat alias — existing code imports `VectorDB` directly.
99	VectorDB = ChromaVectorDB
100		25
101		26
102	# ── Provider selection ────────────────────────────────────────────────────	27	# ── Provider selection ────────────────────────────────────────────────────
⋯ 8 unchanged lines hidden
103	#	28	#
104	# CHAT_VECTOR_STORE_PROVIDER (env var) picks the conversation-memory vector	29	# CHAT_VECTOR_STORE_PROVIDER (env var) picks the conversation-memory vector
105	# store backend. This is intentionally a separate switch from RAG's	30	# store backend. This is intentionally a separate switch from RAG's
106	# VECTOR_STORE_PROVIDER — the two are different indexes by design. Defaults	31	# VECTOR_STORE_PROVIDER — the two are different indexes by design. Defaults
107	# to today's behavior — on-prem/local deployments don't need to set anything.	32	# to today's behavior — on-prem/local deployments don't need to set anything.
108	#	33	#
109	# chroma (default) — local Chroma persistence.	34	# chroma (default) — local Chroma persistence.
110	# azure_search — Azure AI Search, a separate index from the RAG one. Not implemented yet.	35	# azure_search — Azure AI Search, a separate index from the RAG one. Not implemented yet.
111		36
112	def create_conversation_vector_store(	37	def create_conversation_vector_store(
113	collection_name: str,	38	collection_name: str,
		39	embedding: Embeddings,
114	persist_directory: str = "./chroma",	40	persist_directory: str = "./chroma",
115	provider: Optional[str] = None,	41	provider: Optional[str] = None,
116	) -> ConversationVectorStoreBase:	42	) -> VectorStore:
117	"""Factory for the conversation-memory vector store, selected by CHAT_VECTOR_STORE_PROVIDER."""	43	"""Factory for the conversation-memory vector store, selected by CHAT_VECTOR_STORE_PROVIDER."""
118	provider = (provider or os.getenv("CHAT_VECTOR_STORE_PROVIDER", "chroma")).strip().lower()	44	provider = (provider or os.getenv("CHAT_VECTOR_STORE_PROVIDER", "chroma")).strip().lower()
119		45
120	if provider == "chroma":	46	if provider == "chroma":
121	return ChromaVectorDB(collection_name=collection_name, persist_directory=persist_directory)	47	from langchain_chroma import Chroma
		48
		49	return Chroma(
		50	collection_name=collection_name,
		51	embedding_function=embedding,
		52	persist_directory=persist_directory,
		53	# Chroma's product telemetry would otherwise phone home on startup.
		54	client_settings=Settings(anonymized_telemetry=False, is_persistent=True,
		55	persist_directory=persist_directory),
		56	)
122		57
123	if provider == "azure_search":	58	if provider == "azure_search":
124	raise NotImplementedError(	59	raise NotImplementedError(
⋯ 7 unchanged lines hidden
125	"CHAT_VECTOR_STORE_PROVIDER=azure_search is not implemented yet. Use 'chroma' for now."	60	"CHAT_VECTOR_STORE_PROVIDER=azure_search is not implemented yet. Use 'chroma' for now."
126	)	61	)
127		62
128	raise ValueError(	63	raise ValueError(
129	f"Unknown CHAT_VECTOR_STORE_PROVIDER={provider!r}. Supported: chroma, "	64	f"Unknown CHAT_VECTOR_STORE_PROVIDER={provider!r}. Supported: chroma, "
130	"azure_search (coming soon)."	65	"azure_search (coming soon)."
131	)	66	)

19 diagnose.py

scripts/diagnose.py · 48 → 48 lines · 2 rows changed

A diagnostic script. One import changes.

cortex-core · hand-written

Builds OpenAIEmbedder directly.

cortex-langgraph · LangChain / LangGraph

Uses the create_embedder() factory, so it follows EMBEDDING_PROVIDER like the app does.

Why it matters: Scripts and the app now pick providers the same way.

factory

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core scripts/diagnose.py		cortex-langgraph scripts/diagnose.py
⋯ 37 unchanged lines hidden
1	"""Quick diagnostic: check vectorstore contents and embedder work correctly."""	1	"""Quick diagnostic: check vectorstore contents and embedder work correctly."""
2	import sys	2	import sys
3	from pathlib import Path	3	from pathlib import Path
4		4
5	ROOT = Path(__file__).resolve().parents[1]	5	ROOT = Path(__file__).resolve().parents[1]
6	if str(ROOT) not in sys.path:	6	if str(ROOT) not in sys.path:
7	sys.path.insert(0, str(ROOT))	7	sys.path.insert(0, str(ROOT))
8		8
9	from dotenv import load_dotenv	9	from dotenv import load_dotenv
10	import os	10	import os
11		11
12	load_dotenv()	12	load_dotenv()
13		13
14	# 1. Check API key	14	# 1. Check API key
15	api_key = os.getenv("OPENAI_API_KEY")	15	api_key = os.getenv("OPENAI_API_KEY")
16	print(f"1. OPENAI_API_KEY loaded: {'YES' if api_key else 'NO'}")	16	print(f"1. OPENAI_API_KEY loaded: {'YES' if api_key else 'NO'}")
17	if api_key:	17	if api_key:
18	print(f" Key starts with: {api_key[:12]}...")	18	print(f" Key starts with: {api_key[:12]}...")
19		19
20	# 2. Check vectorstore	20	# 2. Check vectorstore
21	persist_dir = str(ROOT / "src" / "business" / "rag" / "vectorstore")	21	persist_dir = str(ROOT / "src" / "business" / "rag" / "vectorstore")
22	print(f"\n2. Vector store path: {persist_dir}")	22	print(f"\n2. Vector store path: {persist_dir}")
23	print(f" Exists: {Path(persist_dir).exists()}")	23	print(f" Exists: {Path(persist_dir).exists()}")
24		24
25	import chromadb	25	import chromadb
26	client = chromadb.PersistentClient(path=persist_dir)	26	client = chromadb.PersistentClient(path=persist_dir)
27	collections = client.list_collections()	27	collections = client.list_collections()
28	print(f" Collections: {[c.name for c in collections]}")	28	print(f" Collections: {[c.name for c in collections]}")
29		29
30	for col in collections:	30	for col in collections:
31	count = col.count()	31	count = col.count()
32	print(f" Collection '{col.name}': {count} records")	32	print(f" Collection '{col.name}': {count} records")
33	if count > 0:	33	if count > 0:
34	peek = col.peek(limit=2)	34	peek = col.peek(limit=2)
35	print(f" Sample IDs: {peek['ids'][:2]}")	35	print(f" Sample IDs: {peek['ids'][:2]}")
36	print(f" Sample doc preview: {peek['documents'][0][:100] if peek['documents'] else 'EMPTY'}...")	36	print(f" Sample doc preview: {peek['documents'][0][:100] if peek['documents'] else 'EMPTY'}...")
37		37
38	# 3. Check embedder	38	# 3. Check embedder
39	print("\n3. Testing embedder...")	39	print("\n3. Testing embedder...")
40	try:	40	try:
41	from src.business.core.embedding import OpenAIEmbedder	41	from src.business.core.embedding import create_embedder
42	embedder = OpenAIEmbedder(api_key=api_key)	42	embedder = create_embedder(api_key=api_key)
43	result = embedder.embed_query("test query")	43	result = embedder.embed_query("test query")
44	print(f" Embedding dimension: {len(result)}")	44	print(f" Embedding dimension: {len(result)}")
45	print(f" First 5 values: {result[:5]}")	45	print(f" First 5 values: {result[:5]}")
⋯ 3 unchanged lines hidden
46	print(" Embedder: OK")	46	print(" Embedder: OK")
47	except Exception as e:	47	except Exception as e:
48	print(f" Embedder ERROR: {e}")	48	print(f" Embedder ERROR: {e}")

20 retrieval_cli.py

scripts/retrieval_cli.py · 17 → 19 lines · 7 rows changed

A command-line check of raw retrieval (no re-ranking, no LLM).

cortex-core · hand-written

Builds the whole RAGPipeline and calls its private _retrieve() method.

cortex-langgraph · LangChain / LangGraph

Calls store.similarity_search_with_score() directly: the standard, public LangChain API.

Why it matters: No reaching into a class's private methods to test one step.

similarity_search_with_score

removed / rewritten in cortex-coreadded / rewritten in cortex-langgraphunchanged

cortex-core scripts/retrieval_cli.py		cortex-langgraph scripts/retrieval_cli.py
⋯ 5 unchanged lines hidden
1	import sys	1	import sys
2	from pathlib import Path	2	from pathlib import Path
3		3
4	# Ensure project root is on sys.path when running directly	4	# Ensure project root is on sys.path when running directly
5	ROOT = Path(__file__).resolve().parents[1]	5	ROOT = Path(__file__).resolve().parents[1]
6	if str(ROOT) not in sys.path:	6	if str(ROOT) not in sys.path:
7	sys.path.insert(0, str(ROOT))	7	sys.path.insert(0, str(ROOT))
8		8
9	from src.business.rag.retrieval import RAGPipeline	9	from src.business.rag.vector_store import create_vector_store
10	from pathlib import Path	10	from pathlib import Path
11	persist_dir = str(Path(__file__).resolve().parents[1] / "src/business/rag/vectorstore")	11	persist_dir = str(Path(__file__).resolve().parents[1] / "src/business/rag/vectorstore")
12	rag = RAGPipeline(persist_dir=persist_dir)	12	store = create_vector_store(persist_dir=persist_dir)
13	query = rag._retrieve("What are the unique total property counts for EPCs per year?", top_k=5)	13	# Retrieval only (no re-rank, no LLM): the vector store's raw top-5.
14	if not query:	14	results = store.similarity_search_with_score("What are the unique total property counts for EPCs per year?", k=5)
		15	if not results:
15	print("No results retrieved")	16	print("No results retrieved")
16	else:	17	else:
17	print(query[0])	18	doc, distance = results[0]
		19	print(f"distance={distance:.4f} metadata={doc.metadata}\n{doc.page_content}")